package server

import io.github.oshai.kotlinlogging.KotlinLogging
import io.modelcontextprotocol.kotlin.sdk.CallToolResult
import io.modelcontextprotocol.kotlin.sdk.ClientCapabilities
import io.modelcontextprotocol.kotlin.sdk.Implementation
import io.modelcontextprotocol.kotlin.sdk.JSONRPCMessage
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
import io.modelcontextprotocol.kotlin.sdk.TextContent
import io.modelcontextprotocol.kotlin.sdk.client.Client
import io.modelcontextprotocol.kotlin.sdk.client.ClientOptions
import io.modelcontextprotocol.kotlin.sdk.server.Server
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
import io.modelcontextprotocol.kotlin.sdk.shared.AbstractTransport
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.currentTime
import kotlinx.coroutines.test.runTest
import kotlinx.coroutines.withTimeout
import org.junit.jupiter.api.Test
import kotlin.test.assertNotNull
import kotlin.test.assertTrue

private val logger = KotlinLogging.logger {}

/**
 * ChannelBasedInMemoryTransport is a transport that uses a Channel to simulate the asynchronous nature of a real
 * network socket.
 * The normal InMemoryTransport is not susceptible to the issue because it uses the client's coroutine for server-side
 * processing of the request. The ChannelBasedInMemoryTransport, as well as StdioServerTransport and likely any other
 * real transport, are susceptible to the issue as client and server are in separate execution contexts.
 */
private class ChannelBasedInMemoryTransport(private val scope: CoroutineScope, private val name: String) :
    AbstractTransport() {
    private lateinit var otherTransport: ChannelBasedInMemoryTransport
    private val incomingChannel = Channel<JSONRPCMessage>(Channel.UNLIMITED)

    companion object {
        fun createLinkedPair(scope: CoroutineScope): Pair<ChannelBasedInMemoryTransport, ChannelBasedInMemoryTransport> {
            val clientTransport = ChannelBasedInMemoryTransport(scope, "ClientTransport")
            val serverTransport = ChannelBasedInMemoryTransport(scope, "ServerTransport")
            clientTransport.otherTransport = serverTransport
            serverTransport.otherTransport = clientTransport
            return Pair(clientTransport, serverTransport)
        }
    }

    override suspend fun start() {
        scope.launch {
            for (message in incomingChannel) {
                logger.debug { "[$name] RECV: $message" }
                _onMessage(message)
            }
        }
    }

    override suspend fun close() {
        if (!incomingChannel.isClosedForReceive) {
            incomingChannel.close()
            otherTransport.incomingChannel.close()
            _onClose()
        }
    }

    override suspend fun send(message: JSONRPCMessage) {
        logger.debug { "[$name] SEND: $message" }
        if (otherTransport.incomingChannel.isClosedForSend) {
            throw IllegalStateException("Not connected")
        }
        otherTransport.incomingChannel.send(message)
    }
}

class ConcurrencyTest {

    @Test
    fun `server blocks concurrent requests due to serial processing`() = runTest {
        val slowToolDelay = 1000L
        // The test timeout must be longer than the slow tool's delay.
        withTimeout(slowToolDelay + 500) {
            // Arrange: Server with tools capability
            val server = Server(
                serverInfo = Implementation("test-server", "1.0"),
                options = ServerOptions(capabilities = ServerCapabilities(tools = ServerCapabilities.Tools(null)))
            )

            // Arrange: Add a "slow" tool and a "fast" tool using the convenient overload
            server.addTool("slow_tool", "A tool that takes a while") {
                logger.info { "[Server] 'slow_tool' handler started." }
                delay(slowToolDelay)
                logger.info { "[Server] 'slow_tool' handler finished." }
                CallToolResult(content = listOf(TextContent("slow_tool_done")))
            }
            server.addTool("fast_tool", "A tool that is quick") {
                logger.info { "[Server] 'fast_tool' handler executed." }
                CallToolResult(content = listOf(TextContent("fast_tool_done")))
            }

            // Arrange: Client
            val client = Client(
                clientInfo = Implementation("test-client", "1.0"),
                options = ClientOptions()
            )

            // Arrange: Transport and connection
            val (clientTransport, serverTransport) = ChannelBasedInMemoryTransport.createLinkedPair(this)
            val serverInitialized = CompletableDeferred<Unit>()
            server.onInitialized { serverInitialized.complete(Unit) }

            launch { server.connect(serverTransport) }
            launch { client.connect(clientTransport) }
            serverInitialized.await()

            // Act: Launch two concurrent tool calls. The slow one first, then the fast one.
            val startTime = currentTime

            val slowJob = launch {
                logger.info { "[Client] Calling 'slow_tool'..." }
                val result = client.callTool("slow_tool", emptyMap())
                assertNotNull(result)
            }
            // Give the first request a moment to be sent and start processing on the server
            delay(50)

            val fastJob = launch {
                logger.info { "[Client] Calling 'fast_tool'..." }
                val result = client.callTool("fast_tool", emptyMap())
                assertNotNull(result)
            }

            // Wait for both to complete
            slowJob.join()
            fastJob.join()

            val duration = currentTime - startTime

            // Assert: The total time should be greater than the slow tool's delay,
            // proving the fast call was blocked by the slow one. If they ran in parallel,
            // the total duration would be only slightly more than the slow tool's delay.
            // Because they run serially, the duration is > slow + fast + overhead.
            logger.info { "Total execution time: ${duration}ms" }
            assertTrue(
                duration >= slowToolDelay,
                "Fast tool was blocked by slow tool. Total duration (${duration}ms) should be >= slow tool delay (${slowToolDelay}ms)."
            )
        }
    }
}