/*
 * Copyright 2026 Sergey S. Chernov real.sergeych@gmail.com
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package net.sergeych.lyng.obj

import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonObject
import net.sergeych.lyng.CallSignature
import net.sergeych.lyng.Scope
import net.sergeych.lyng.miniast.*
import net.sergeych.lynon.LynonDecoder
import net.sergeych.lynon.LynonEncoder
import net.sergeych.lynon.LynonType

class ObjMapEntry(val key: Obj, val value: Obj) : Obj() {

    override suspend fun compareTo(scope: Scope, other: Obj): Int {
        if (other !is ObjMapEntry) return -1
        val c = key.compareTo(scope, other.key)
        if (c != 0) return c
        return value.compareTo(scope, other.value)
    }

    override fun hashCode(): Int {
        return key.hashCode() + value.hashCode()
    }

    override fun equals(other: Any?): Boolean {
        return other is ObjMapEntry && key == other.key && value == other.value
    }

    override suspend fun getAt(scope: Scope, index: Obj): Obj = when (index.toInt()) {
        0 -> key
        1 -> value
        else -> scope.raiseIndexOutOfBounds()
    }

    override suspend fun defaultToString(scope: Scope): ObjString {
        return ObjString("(${key.toString(scope).value} => ${value.toString(scope).value})")
    }

    override val objClass get() = type

    override suspend fun serialize(scope: Scope, encoder: LynonEncoder, lynonType: LynonType?) {
        encoder.encodeAny(scope,key)
        encoder.encodeAny(scope,value)
    }

    companion object {
        val type = object : ObjClass("MapEntry", ObjArray) {
            override suspend fun callOn(scope: Scope): Obj {
                return ObjMapEntry(scope.requiredArg<Obj>(0), scope.requiredArg<Obj>(1))
            }

            override suspend fun deserialize(scope: Scope, decoder: LynonDecoder, lynonType: LynonType?): Obj {
                return ObjMapEntry(
                    decoder.decodeAny(scope),
                    decoder.decodeAny(scope)
                )
            }
        }.apply {
            addPropertyDoc(
                name = "key",
                doc = "Key component of this map entry.",
                type = type("lyng.Any"),
                moduleName = "lyng.stdlib",
                getter = { thisAs<ObjMapEntry>().key }
            )
            addPropertyDoc(
                name = "value",
                doc = "Value component of this map entry.",
                type = type("lyng.Any"),
                moduleName = "lyng.stdlib",
                getter = { thisAs<ObjMapEntry>().value }
            )
            addPropertyDoc(
                name = "size",
                doc = "Number of components in this entry (always 2).",
                type = type("lyng.Int"),
                moduleName = "lyng.stdlib",
                getter = { 2.toObj() }
            )
        }
    }

    override suspend fun plus(scope: Scope, other: Obj): Obj {
        // Build a new map starting from this entry, then merge `other`.
        val result = ObjMap(mutableMapOf(key to value))
        return result.plus(scope, other)
    }
}

class ObjMap(val map: MutableMap<Obj, Obj> = mutableMapOf()) : Obj() {

    override suspend fun equals(scope: Scope, other: Obj): Boolean {
        if (this === other) return true
        val otherSize = when (other) {
            is ObjMap -> other.map.size
            is ObjImmutableMap -> other.map.size
            else -> return false
        }
        if (map.size != otherSize) return false
        for ((k, v) in map) {
            val otherV = other.getAt(scope, k)
            if (otherV === ObjNull && !other.contains(scope, k)) return false
            if (!v.equals(scope, otherV)) return false
        }
        return true
    }

    override val objClass get() = type

    override suspend fun getAt(scope: Scope, index: Obj): Obj =
        map.get(index) ?: ObjNull

    override suspend fun putAt(scope: Scope, index: Obj, newValue: Obj) {
        map[index] = newValue
    }

    override suspend fun contains(scope: Scope, other: Obj): Boolean {
        return other in map
    }

    override suspend fun compareTo(scope: Scope, other: Obj): Int {
        val otherMap = when (other) {
            is ObjMap -> other.map
            is ObjImmutableMap -> other.map
            else -> return -1
        }
        if (map == otherMap) return 0
        if (map.size != otherMap.size) return map.size.compareTo(otherMap.size)
        // for same size, if they are not equal, we don't have a stable order
        // but let's try to be consistent
        return map.toString().compareTo(otherMap.toString())
    }

    override suspend fun defaultToString(scope: Scope): ObjString {
        val reusult = buildString {
            append("Map(")
            var first = true
            for( (k,v) in map) {
                if( !first ) append(",")
                append(k.inspect(scope))
                append(" => ")
                append(v.toString(scope).value)
                first = false
            }
            append(")")
        }
        return ObjString(reusult)
    }

    override fun hashCode(): Int {
        return map.hashCode()
    }

    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (other == null || this::class != other::class) return false

        other as ObjMap

        return map == other.map
    }

    override suspend fun lynonType(): LynonType = LynonType.Map

    override suspend fun serialize(scope: Scope, encoder: LynonEncoder, lynonType: LynonType?) {
        val keys = map.keys.map { it.toObj() }
        val values = map.values.map { it.toObj() }
        encoder.encodeAnyList(scope, keys)
        encoder.encodeAnyList(scope, values, fixedSize = true)
    }

    override suspend fun toJson(scope: Scope): JsonElement {
        return JsonObject(
            map.map { it.key.toString(scope).value to it.value.toJson(scope) }.toMap()
        )
    }

    companion object {

        suspend fun listToMap(scope: Scope, list: List<Obj>): MutableMap<Obj, Obj> {
            val map = mutableMapOf<Obj, Obj>()
            if (list.isEmpty()) return map

            list.forEach {
                when (it) {
                    is ObjMapEntry -> map[it.key] = it.value
                    else -> {
                        if (it.isInstanceOf(ObjArray)) {
                            if (it.invokeInstanceMethod(scope, "size").toInt() != 2)
                                scope.raiseIllegalArgument(
                                    "Array to construct map entry should exactly be 2 elements [key,value], got $it"
                                )
                            map[it.getAt(scope, ObjInt.Zero)] = it.getAt(scope, ObjInt.One)
                        } else {
                            scope.raiseIllegalArgument("elements to construct map must be MapEntry or Array of 2 elements; got $it")
                        }
                    }
                }
            }
            return map
        }

        val type = object : ObjClass("Map", ObjCollection) {
            override suspend fun callOn(scope: Scope): Obj {
                return ObjMap(listToMap(scope, scope.args.list))
            }

            override suspend fun deserialize(scope: Scope, decoder: LynonDecoder, lynonType: LynonType?): Obj {
                val keys = decoder.decodeAnyList(scope)
                val values = decoder.decodeAnyList(scope,fixedSize = keys.size)
                if( keys.size != values.size) scope.raiseIllegalArgument("map keys and values should be same size")
                return ObjMap(keys.zip(values).toMap().toMutableMap())
            }
        }.apply {
            implementingNames.add("Delegate")
            addFn("getValue") {
                val self = thisAs<ObjMap>()
                val key = requiredArg<Obj>(1)
                self.map[key] ?: ObjNull
            }
            addFn("setValue") {
                val self = thisAs<ObjMap>()
                val key = requiredArg<Obj>(1)
                val value = requiredArg<Obj>(2)
                self.map[key] = value
                self
            }
            addFn("bind") {
                val mode = requiredArg<ObjEnumEntry>(1)
                if(  mode.ordinal.value > 1)
                    raiseIllegalArgument("Map can be delegated only to val or var, got ${mode.name.value}")
                thisObj
            }
            addFnDoc(
                name = "getOrNull",
                doc = "Get value by key or return null if the key is absent.",
                params = listOf(ParamDoc("key")),
                returns = type("lyng.Any", nullable = true),
                moduleName = "lyng.stdlib"
            ) {
                val key = args.firstAndOnly(pos)
                thisAs<ObjMap>().map.getOrElse(key) { ObjNull }
            }
            addFnDoc(
                name = "getOrPut",
                doc = "Get value by key or compute, store, and return the default from a lambda.",
                params = listOf(ParamDoc("key"), ParamDoc("default")),
                returns = type("lyng.Any"),
                moduleName = "lyng.stdlib",
                callSignature = CallSignature(
                    inlineHigherOrder = CallSignature.HigherOrderInline(
                        kind = CallSignature.Kind.MAP_GET_OR_PUT,
                        result = CallSignature.ResultMode.BLOCK_RESULT,
                        argCount = 2,
                        lambdaArgIndex = 1
                    )
                )
            ) {
                val key = requiredArg<Obj>(0)
                thisAs<ObjMap>().map.getOrPut(key) {
                    val lambda = requiredArg<Obj>(1)
                    call(lambda)
                }
            }
            addPropertyDoc(
                name = "size",
                doc = "Number of entries in the map.",
                type = type("lyng.Int"),
                moduleName = "lyng.stdlib",
                getter = { (this.thisObj as ObjMap).map.size.toObj() }
            )
            addFnDoc(
                name = "remove",
                doc = "Remove the entry by key and return the previous value or null if absent.",
                params = listOf(ParamDoc("key")),
                returns = type("lyng.Any", nullable = true),
                moduleName = "lyng.stdlib"
            ) {
                thisAs<ObjMap>().map.remove(requiredArg<Obj>(0))?.toObj() ?: ObjNull
            }
            addFnDoc(
                name = "clear",
                doc = "Remove all entries from this map. Returns the map.",
                returns = type("lyng.Map"),
                moduleName = "lyng.stdlib"
            ) {
                thisAs<ObjMap>().map.clear()
                thisObj
            }
            addPropertyDoc(
                name = "keys",
                doc = "List of keys in this map.",
                type = TypeGenericDoc(type("lyng.List"), listOf(type("lyng.Any"))),
                moduleName = "lyng.stdlib",
                getter = { thisAs<ObjMap>().map.keys.toObj() }
            )
            addPropertyDoc(
                name = "values",
                doc = "List of values in this map.",
                type = TypeGenericDoc(type("lyng.List"), listOf(type("lyng.Any"))),
                moduleName = "lyng.stdlib",
                getter = { ObjList(thisAs<ObjMap>().map.values.toMutableList()) }
            )
            addFnDoc(
                name = "iterator",
                doc = "Iterator over map entries as MapEntry objects.",
                returns = TypeGenericDoc(type("lyng.Iterator"), listOf(type("lyng.MapEntry"))),
                moduleName = "lyng.stdlib"
            ) {
                ObjKotlinIterator(thisAs<ObjMap>().map.entries.iterator())
            }
            addFnDoc(
                name = "toImmutable",
                doc = "Create an immutable snapshot of this map.",
                returns = type("lyng.ImmutableMap"),
                moduleName = "lyng.stdlib"
            ) {
                ObjImmutableMap(thisAs<ObjMap>().map)
            }
        }
    }

    // Merge operations
    override suspend fun plus(scope: Scope, other: Obj): Obj {
        val result = ObjMap(map.toMutableMap())
        result.mergeIn(scope, other)
        return result
    }

    override suspend fun plusAssign(scope: Scope, other: Obj): Obj {
        mergeIn(scope, other)
        return this
    }

    private suspend fun mergeIn(scope: Scope, other: Obj) {
        when (other) {
            is ObjMap -> {
                // Rightmost wins: copy all entries from `other` over existing ones
                for ((k, v) in other.map) {
                    map[k] = v
                }
            }
            is ObjImmutableMap -> {
                for ((k, v) in other.map) {
                    map[k] = v
                }
            }
            is ObjMapEntry -> {
                map[other.key] = other.value
            }
            is ObjList -> {
                // Treat as list of map entries
                for (e in other.list) {
                    val entry = when (e) {
                        is ObjMapEntry -> e
                        else -> {
                            if (e.isInstanceOf(ObjArray)) {
                                if (e.invokeInstanceMethod(scope, "size").toInt() != 2)
                                    scope.raiseIllegalArgument("Array element to merge into map must have 2 elements, got $e")
                                ObjMapEntry(e.getAt(scope, 0), e.getAt(scope, 1))
                            } else
                                scope.raiseIllegalArgument("map can only be merged with MapEntry elements; got $e")
                        }
                    }
                    map[entry.key] = entry.value
                }
            }
            else -> scope.raiseIllegalArgument("map can only be merged with Map, MapEntry, or List<MapEntry>")
        }
    }
}
