"""Every wait on the MCP path is bounded.

A server that accepts a request and never answers is the shape of the prod
failure. With no read timeout the SDK waits forever (``anyio.fail_after(None)``),
the caller's future is never resolved, and the runner task stays wedged for the
life of the process — so every later reload blocked on that same client and
``disconnect_all`` walked the pool one stuck client at a time.

These tests talk to a fake transport through a *real* ``ClientSession``: it reads
requests and replies to nothing, which is what a silent server looks like from
here. The bounds are the code under test, so each test patches them down to
fractions of a second.
"""

import asyncio
import contextlib
import time

import anyio
import httpx
import pytest
from mcp.shared.message import SessionMessage
from mcp.types import (
    LATEST_PROTOCOL_VERSION,
    Implementation,
    InitializeResult,
    JSONRPCMessage,
    JSONRPCResponse,
    ServerCapabilities,
)

import navi.mcp.client as client_mod
from navi.mcp.client import McpClient, McpOperationTimeout
from navi.mcp.config import McpServerConfig
from navi.mcp.manager import McpManager


class SilentServer:
    """A stand-in stdio transport: reads every request, answers almost none.

    ``answer_handshake=False`` makes it silent from the very first byte — the
    server that accepts the connection and says nothing at all. Otherwise it
    completes the ``initialize`` handshake (so the session is real and usable)
    and then goes quiet, which is how a server that died mid-session behaves.
    """

    def __init__(self, *, answer_handshake: bool = True) -> None:
        self.answer_handshake = answer_handshake
        self.requests: list[str] = []
        self._tasks: list[asyncio.Task] = []

    def __call__(self, params):  # stands in for stdio_client(params)
        return self

    async def __aenter__(self):
        # (read_stream, write_stream) from the client's side: it reads what we
        # send, and we read what it writes. Streams are buffered, so a client
        # write never blocks on us reading it.
        to_client_w, to_client = anyio.create_memory_object_stream[SessionMessage](8)
        from_client, from_client_r = anyio.create_memory_object_stream[SessionMessage](8)
        self._to_client = to_client_w
        self._from_client = from_client
        self._tasks.append(asyncio.create_task(self._serve(from_client_r)))
        return to_client, from_client

    async def __aexit__(self, *exc_info) -> bool:
        for task in self._tasks:
            task.cancel()
        for task in self._tasks:
            with contextlib.suppress(BaseException):
                await task
        self._tasks.clear()
        return False

    async def _serve(self, reader) -> None:
        async for session_message in reader:
            request = session_message.message.root
            method = getattr(request, "method", None)
            if method is None:
                continue
            self.requests.append(method)
            if method != "initialize" or not self.answer_handshake:
                continue  # read it, never answer it
            result = InitializeResult(
                protocolVersion=LATEST_PROTOCOL_VERSION,
                capabilities=ServerCapabilities(),
                serverInfo=Implementation(name="silent", version="1"),
            )
            await self._to_client.send(
                SessionMessage(
                    message=JSONRPCMessage(
                        JSONRPCResponse(
                            jsonrpc="2.0",
                            id=request.id,
                            result=result.model_dump(by_alias=True, mode="json"),
                        )
                    )
                )
            )


def _cfg() -> McpServerConfig:
    return McpServerConfig(transport="stdio", command="python", args=["-c", "pass"])


@pytest.fixture
def quiet_after_handshake(monkeypatch):
    """A server that completes the handshake, then never answers again."""
    server = SilentServer(answer_handshake=True)
    monkeypatch.setattr(client_mod, "stdio_client", server)
    return server


@pytest.fixture
def silent_from_the_start(monkeypatch):
    """A server that accepts the connection and answers nothing at all."""
    server = SilentServer(answer_handshake=False)
    monkeypatch.setattr(client_mod, "stdio_client", server)
    return server


async def test_a_silent_server_cannot_wedge_list_tools(quiet_after_handshake, monkeypatch):
    """The SDK-level read timeout is what ends a request nobody answers — and it
    ends it *inside* the runner, so the client is still usable afterwards."""
    monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 0.2)
    monkeypatch.setattr(client_mod, "_LIST_TOOLS_TIMEOUT_S", 10.0)
    client = McpClient("quiet", _cfg())
    await asyncio.wait_for(client.connect(), timeout=5)

    started = time.monotonic()
    with pytest.raises(Exception, match="Timed out"):
        await asyncio.wait_for(client.list_tools(), timeout=5)

    assert time.monotonic() - started < 3
    assert quiet_after_handshake.requests.count("tools/list") >= 1
    await client.disconnect()


async def test_the_caller_bound_covers_an_op_the_sdk_never_times_out(
    quiet_after_handshake, monkeypatch
):
    """The layer below the SDK's: an op that never reaches the SDK's read wait —
    a server that stops reading its requests, a caller queued behind one — is
    caught here, and the wedged runner is aborted rather than left for the next
    command to queue behind."""
    monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 60.0)
    monkeypatch.setattr(client_mod, "_LIST_TOOLS_TIMEOUT_S", 0.2)
    monkeypatch.setattr(client_mod, "_CONNECT_TIMEOUT_S", 5.0)
    client = McpClient("wedged", _cfg())
    await client.connect()

    with pytest.raises(McpOperationTimeout) as excinfo:
        await asyncio.wait_for(client.list_tools(), timeout=5)

    assert excinfo.value.server == "wedged"
    assert excinfo.value.op == "list_tools"
    # Aborted, not merely abandoned: the next command starts a fresh runner.
    assert client._runner_task is None
    assert client.connected is False


async def test_a_server_that_never_answers_cannot_hang_load_all(
    silent_from_the_start, monkeypatch
):
    """Startup (and every reload) connects to every configured server. One that
    accepts the connection and says nothing used to hang the whole loop."""
    monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 0.2)
    manager = McpManager()

    started = time.monotonic()
    await asyncio.wait_for(
        manager.load_all({"quiet": _cfg()}), timeout=5
    )

    assert time.monotonic() - started < 3
    assert "quiet" in manager.clients  # kept in the pool for the health check
    assert manager._connected_status["quiet"] is False
    await manager.disconnect_all()


async def test_a_wedged_client_does_not_block_load_all(silent_from_the_start, monkeypatch):
    """The prod symptom, at the manager level: ``load_all`` disconnects the whole
    pool first, so one client stuck in teardown blocked every reload after it."""
    monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 60.0)
    monkeypatch.setattr(client_mod, "_CONNECT_TIMEOUT_S", 60.0)
    monkeypatch.setattr(client_mod, "_DISCONNECT_TIMEOUT_S", 0.2)

    manager = McpManager()
    wedged = McpClient("wedged", _cfg())
    manager._clients["wedged"] = wedged
    pending = asyncio.create_task(wedged.connect())  # never answers
    await asyncio.sleep(0.1)

    started = time.monotonic()
    await asyncio.wait_for(manager.load_all({}), timeout=5)

    assert time.monotonic() - started < 3
    assert manager.clients == {}

    # The caller that was waiting on the abandoned runner is released, not left
    # hanging: a reload's budget has to be able to reach the end of its work.
    with pytest.raises(BaseException):
        await asyncio.wait_for(pending, timeout=2)


async def test_the_streamable_http_transport_is_not_left_on_httpx_defaults(monkeypatch):
    """A bare ``httpx.AsyncClient`` defaults to 5s for *everything*, so the GET
    stream of an idle streamable-http server — and the POST that carries a slow
    tool call — died on a timeout nobody chose."""
    monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 0.2)
    captured: dict = {}

    class RecordingClient(httpx.AsyncClient):
        def __init__(self, **kwargs):
            captured.update(kwargs)
            super().__init__(**kwargs)

    monkeypatch.setattr(client_mod.httpx, "AsyncClient", RecordingClient)
    client = McpClient(
        "streamy",
        McpServerConfig(transport="streamable_http", url="http://127.0.0.1:9/mcp"),
    )
    with pytest.raises(Exception):
        await client.connect()

    assert captured["timeout"].read == client_mod._HTTP_STREAM_READ_TIMEOUT_S
    assert captured["timeout"].connect == client_mod._HTTP_TIMEOUT_S
    await client.disconnect()