package org.ronsmits.nestor.mcp

import ai.koog.agents.core.tools.ToolDescriptor
import ai.koog.agents.mcp.DefaultMcpToolDescriptorParser
import ai.koog.agents.mcp.McpToolDescriptorParser
import kotlinx.serialization.json.*
import org.slf4j.LoggerFactory
import io.modelcontextprotocol.kotlin.sdk.types.Tool as SDKTool

/**
 * Wraps Koog's [DefaultMcpToolDescriptorParser] with JSON Schema normalization
 * that handles patterns the default parser chokes on:
 *
 * 1. `type` as an array (`["string", "null"]`) — picks the first non-null type.
 * 2. Missing `type`/`anyOf`/`enum` (untyped params) — defaults to `"string"`.
 *
 * Normalization is applied recursively to nested object properties and array items.
 */
object LenientMcpToolDescriptorParser : McpToolDescriptorParser {

    private val log = LoggerFactory.getLogger("LenientMcpParser")

    override fun parse(sdkTool: SDKTool): ToolDescriptor {
        val original = sdkTool.inputSchema.properties
        if (original == null) return DefaultMcpToolDescriptorParser.parse(sdkTool)

        val normalized = JsonObject(original.mapValues { (name, element) ->
            if (element is JsonObject) normalizeProperty(name, element) else element
        })

        val fixedTool = if (normalized != original) {
            log.debug("Normalized JSON Schema for tool '{}'", sdkTool.name)
            sdkTool.copy(inputSchema = sdkTool.inputSchema.copy(properties = normalized))
        } else {
            sdkTool
        }

        return DefaultMcpToolDescriptorParser.parse(fixedTool)
    }

    private fun normalizeProperty(path: String, obj: JsonObject): JsonObject {
        val map = obj.toMutableMap()
        var changed = false

        val typeEl = map["type"]
        if (typeEl != null && typeEl !is JsonPrimitive) {
            try {
                val types = typeEl.jsonArray.map { it.jsonPrimitive.content }
                val nonNull = types.firstOrNull { it != "null" } ?: "string"
                map["type"] = JsonPrimitive(nonNull)
                changed = true
                log.debug("Schema fix [{}]: type {} → \"{}\"", path, types, nonNull)
            } catch (e: Exception) {
                log.warn("Schema fix [{}]: could not normalize type array: {}", path, e.message)
            }
        }

        if ("type" !in map && "anyOf" !in map && "enum" !in map) {
            map["type"] = JsonPrimitive("string")
            changed = true
            log.debug("Schema fix [{}]: added default type \"string\" to untyped param", path)
        }

        val items = map["items"]
        if (items is JsonObject) {
            val fixedItems = normalizeProperty("$path.items", items)
            if (fixedItems !== items) { map["items"] = fixedItems; changed = true }
        }

        val props = map["properties"]
        if (props is JsonObject) {
            val fixedProps = JsonObject(props.mapValues { (n, v) ->
                if (v is JsonObject) normalizeProperty("$path.$n", v) else v
            })
            if (fixedProps != props) { map["properties"] = fixedProps; changed = true }
        }

        return if (changed) JsonObject(map) else obj
    }
}