package com.correx.infrastructure.tools.mcp

import com.correx.core.tools.process.ChildProcess
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonArray
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonNull
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.decodeFromJsonElement
import kotlinx.serialization.json.put
import kotlinx.serialization.json.putJsonObject
import java.io.BufferedReader
import java.io.Writer
import java.util.concurrent.atomic.AtomicLong

/**
 * The transport an [McpClient] talks over. Extracted so tests can drive the client with an
 * in-process fake instead of spawning a real subprocess (keeps the client's protocol logic
 * deterministically testable — no process, no I/O).
 */
internal interface McpTransport {
    suspend fun request(request: JsonRpcRequest): JsonRpcResponse
    suspend fun notify(notification: JsonRpcNotification)
    fun close()
}

private val MCP_JSON = Json {
    ignoreUnknownKeys = true
    encodeDefaults = true
}

/**
 * A live MCP session: performs the `initialize` handshake, lists tools, and calls them. Ecosystem-
 * agnostic — it neither knows nor cares whether the server behind [transport] resolves npm packages,
 * indexes a codebase, or anything else. The nondeterminism of a `tools/call` is captured by the
 * kernel's normal tool-execution recording (the call rides the ToolExecutor path), so nothing here
 * needs to emit events: replay never re-calls because replay never executes tools.
 */
internal class McpClient(
    val serverId: String,
    private val transport: McpTransport,
) : McpClientHandle {
    suspend fun initialize() {
        val params = buildJsonObject {
            put("protocolVersion", MCP_PROTOCOL_VERSION)
            putJsonObject("capabilities") {}
            putJsonObject("clientInfo") {
                put("name", "correx")
                put("version", "0")
            }
        }
        transport.request(JsonRpcRequest(method = "initialize", params = params, id = 0)).orThrow(serverId)
        transport.notify(JsonRpcNotification(method = "notifications/initialized"))
    }

    suspend fun listTools(): List<McpToolDescriptor> {
        val result = transport.request(JsonRpcRequest(method = "tools/list", id = nextId())).orThrow(serverId)
        return MCP_JSON.decodeFromJsonElement(McpListToolsResult.serializer(), result).tools
    }

    override suspend fun callTool(remoteName: String, arguments: JsonObject): McpToolCallOutcome {
        val params = buildJsonObject {
            put("name", remoteName)
            put("arguments", arguments)
        }
        val result = transport.request(JsonRpcRequest(method = "tools/call", params = params, id = nextId()))
            .orThrow(serverId)
        val decoded = MCP_JSON.decodeFromJsonElement(McpCallToolResult.serializer(), result)
        return McpToolCallOutcome(text = decoded.text(), isError = decoded.isError)
    }

    fun close() = transport.close()

    private val ids = AtomicLong(1)
    private fun nextId() = ids.getAndIncrement()

    private fun JsonRpcResponse.orThrow(serverId: String): JsonElement {
        error?.let { error("MCP server '$serverId' error ${it.code}: ${it.message}") }
        return result ?: error("MCP server '$serverId' returned no result")
    }

    companion object {
        // The protocol revision Correx advertises. Servers negotiate down if they only speak older.
        const val MCP_PROTOCOL_VERSION = "2024-11-05"

        /**
         * Spawn [command] as a child process and return an initialized client speaking MCP over its
         * stdio pipes. The caller owns [McpClient.close] (kills the process). Uses [ChildProcess] so
         * the server inherits the same unattended/non-interactive execution guarantees as every other
         * spawned child.
         */
        suspend fun spawn(serverId: String, command: List<String>, env: Map<String, String>): McpClient {
            val process = withContext(Dispatchers.IO) {
                ChildProcess.builder(command)
                    .apply {
                        environment().putAll(env)
                        // ChildProcess defaults stdin to /dev/null (so scaffolders fail fast instead of
                        // hanging). MCP stdio is the opposite: the server READS requests from stdin, so
                        // it needs a live pipe or it sees EOF and exits immediately. Restore the pipe.
                        redirectInput(ProcessBuilder.Redirect.PIPE)
                    }
                    .start()
            }
            val client = McpClient(serverId, StdioTransport(process))
            client.initialize()
            return client
        }
    }
}

/**
 * Newline-delimited JSON-RPC over a child process's stdio (the MCP stdio transport). Requests are
 * serialized through a [Mutex]: one line written, then response lines drained until the matching id
 * arrives — server-initiated notifications/logs on the same pipe are skipped, not mistaken for the
 * reply. Blocking pipe I/O is confined to [Dispatchers.IO].
 */
private class StdioTransport(private val process: Process) : McpTransport {
    private val lock = Mutex()
    private val writer: Writer = process.outputStream.bufferedWriter()
    private val reader: BufferedReader = process.inputStream.bufferedReader()

    override suspend fun request(request: JsonRpcRequest): JsonRpcResponse = lock.withLock {
        withContext(Dispatchers.IO) {
            writeLine(MCP_JSON.encodeToString(JsonRpcRequest.serializer(), request))
            while (true) {
                val line = reader.readLine()
                    ?: error("MCP server '${request.method}' closed the connection")
                if (line.isBlank()) continue
                val response = runCatching {
                    MCP_JSON.decodeFromString(JsonRpcResponse.serializer(), line)
                }.getOrNull() ?: continue // not a response we can parse (a notification/log line) — skip
                if (response.id == request.id) return@withContext response
            }
            @Suppress("UNREACHABLE_CODE")
            error("unreachable")
        }
    }

    override suspend fun notify(notification: JsonRpcNotification) = lock.withLock {
        withContext(Dispatchers.IO) {
            writeLine(MCP_JSON.encodeToString(JsonRpcNotification.serializer(), notification))
        }
    }

    override fun close() {
        runCatching { writer.close() }
        runCatching { reader.close() }
        process.destroy()
    }

    private fun writeLine(json: String) {
        writer.write(json)
        writer.write("\n")
        writer.flush()
    }
}

/** Encode a [ToolRequest]'s parameter map into the JsonObject `arguments` payload for a tools/call. */
internal fun argumentsOf(parameters: Map<String, Any?>): JsonObject =
    JsonObject(parameters.mapValues { (_, value) -> anyToJsonElement(value) })

private fun anyToJsonElement(value: Any?): JsonElement = when (value) {
    null -> JsonNull
    is JsonElement -> value
    is Boolean -> JsonPrimitive(value)
    is Number -> JsonPrimitive(value)
    is String -> JsonPrimitive(value)
    is Map<*, *> -> JsonObject(value.entries.associate { (k, v) -> k.toString() to anyToJsonElement(v) })
    is Iterable<*> -> JsonArray(value.map { anyToJsonElement(it) })
    else -> JsonPrimitive(value.toString())
}
