package com.correx.infrastructure.tools.mcp

import com.correx.core.approvals.Tier
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.contract.ToolResult
import kotlinx.coroutines.runBlocking
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.put
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

/**
 * Drives [McpClient] through an in-process [McpTransport] — no subprocess, no I/O. Verifies the
 * protocol logic (handshake, tools/list decode, tools/call outcome, error surfacing) and the
 * [McpTool] adapter's mapping onto the Correx tool contract, all deterministically.
 */
class McpClientTest {

    /** A scripted transport: answers each method from a canned map; records notifications. */
    private class FakeTransport(private val results: Map<String, JsonRpcResponse>) : McpTransport {
        val notified = mutableListOf<String>()
        override suspend fun request(request: JsonRpcRequest): JsonRpcResponse =
            (results[request.method] ?: error("no scripted result for ${request.method}"))
                .copy(id = request.id)
        override suspend fun notify(notification: JsonRpcNotification) { notified.add(notification.method) }
        override fun close() {}
    }

    private fun ok(json: kotlinx.serialization.json.JsonElement) = JsonRpcResponse(result = json)

    @Test
    fun `initialize handshakes then lists and calls tools`(): Unit = runBlocking {
        val listResult = buildJsonObject {
            put("tools", kotlinx.serialization.json.buildJsonArray {
                add(buildJsonObject {
                    put("name", "search_code")
                    put("description", "search the graph")
                })
            })
        }
        val callResult = buildJsonObject {
            put("content", kotlinx.serialization.json.buildJsonArray {
                add(buildJsonObject { put("type", "text"); put("text", "hit one") })
                add(buildJsonObject { put("type", "text"); put("text", "hit two") })
            })
        }
        val transport = FakeTransport(
            mapOf(
                "initialize" to ok(buildJsonObject {}),
                "tools/list" to ok(listResult),
                "tools/call" to ok(callResult),
            ),
        )
        val client = McpClient("cm", transport)
        client.initialize()
        assertEquals(listOf("notifications/initialized"), transport.notified)

        val tools = client.listTools()
        assertEquals(listOf("search_code"), tools.map { it.name })

        val outcome = client.callTool("search_code", JsonObject(emptyMap()))
        assertEquals("hit one\nhit two", outcome.text)
        assertTrue(!outcome.isError)
    }

    @Test
    fun `server error is surfaced as a thrown failure`(): Unit = runBlocking {
        val transport = FakeTransport(
            mapOf("tools/call" to JsonRpcResponse(error = JsonRpcError(code = -32000, message = "boom"))),
        )
        val client = McpClient("cm", transport)
        val ex = runCatching { client.callTool("x", JsonObject(emptyMap())) }.exceptionOrNull()
        assertTrue(ex?.message?.contains("boom") == true, "expected server error message, got: ${ex?.message}")
    }

    @Test
    fun `McpTool maps an isError result to a recoverable failure`(): Unit = runBlocking {
        val handle = object : McpClientHandle {
            override suspend fun callTool(remoteName: String, arguments: JsonObject) =
                McpToolCallOutcome(text = "package not found", isError = true)
        }
        val tool = McpTool("cm", "resolve", "", JsonObject(emptyMap()), Tier.T2, handle)
        assertEquals("mcp__cm__resolve", tool.name)
        val result = tool.execute(req(tool.name)) as ToolResult.Failure
        assertEquals("package not found", result.reason)
        assertTrue(result.recoverable)
    }

    @Test
    fun `McpTool maps a normal result to success`(): Unit = runBlocking {
        val handle = object : McpClientHandle {
            override suspend fun callTool(remoteName: String, arguments: JsonObject) =
                McpToolCallOutcome(text = "ok", isError = false)
        }
        val tool = McpTool("cm", "resolve", "", JsonObject(emptyMap()), Tier.T2, handle)
        val result = tool.execute(req(tool.name)) as ToolResult.Success
        assertEquals("ok", result.output)
    }

    private fun req(name: String) = ToolRequest(
        invocationId = ToolInvocationId("inv-1"),
        sessionId = SessionId("s-1"),
        stageId = StageId("st-1"),
        toolName = name,
        parameters = mapOf("q" to "vite"),
    )
}
