"""Tests for `ServerSession`.

`ServerSession` is a thin per-request proxy over two `Outbound` channels and a
`Connection`. Tested with stub outbounds so we can assert what reaches the wire
(method, params, `CallOptions`) and which channel it routed to, without standing
up a transport.
"""

from collections.abc import Mapping
from typing import Any

import pytest
from pydantic import ValidationError

from mcp import types
from mcp.server.connection import Connection
from mcp.server.session import ServerSession
from mcp.shared.dispatcher import CallOptions
from mcp.shared.message import ServerMessageMetadata
from mcp.shared.version import LATEST_HANDSHAKE_VERSION, LATEST_MODERN_VERSION
from mcp.types import (
    ClientCapabilities,
    Implementation,
    SamplingCapability,
    SamplingToolsCapability,
)


class StubOutbound:
    """Records `send_raw_request` / `notify` / `progress` calls and returns a canned result.

    Structurally a `DispatchContext[Any]` so it can stand in for the per-request channel.
    """

    transport: Any = None
    can_send_request: bool = True
    request_id: Any = None
    message_metadata: Any = None
    cancel_requested: Any = None

    def __init__(self, result: dict[str, Any] | None = None) -> None:
        self.requests: list[tuple[str, Mapping[str, Any] | None, CallOptions | None]] = []
        self.notifications: list[tuple[str, Mapping[str, Any] | None]] = []
        self.progress_calls: list[tuple[float, float | None, str | None]] = []
        self.result = result if result is not None else {}

    async def send_raw_request(
        self,
        method: str,
        params: Mapping[str, Any] | None,
        opts: CallOptions | None = None,
    ) -> dict[str, Any]:
        self.requests.append((method, params, opts))
        return self.result

    async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None:
        self.notifications.append((method, params))

    async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
        self.progress_calls.append((progress, total, message))


def _make_session(
    outbound: StubOutbound,
    *,
    capabilities: ClientCapabilities | None = None,
    protocol_version: str = LATEST_HANDSHAKE_VERSION,
) -> ServerSession:
    """Single-channel session: the stub is both request and standalone outbound."""
    client_info = Implementation(name="c", version="0") if capabilities is not None else None
    conn = Connection.from_envelope(protocol_version, client_info, capabilities, outbound=outbound)
    return ServerSession(outbound, conn)


def _two_channel_session(request_ch: StubOutbound, standalone_ch: StubOutbound) -> ServerSession:
    """Distinct request/standalone outbounds so routing assertions can tell the channels apart."""
    conn = Connection.from_envelope(LATEST_HANDSHAKE_VERSION, None, None, outbound=standalone_ch)
    return ServerSession(request_ch, conn)


@pytest.mark.anyio
async def test_send_request_forwards_timeout_and_progress_callback_as_call_options():
    outbound = StubOutbound(result={"roots": []})
    session = _make_session(outbound)

    async def on_progress(progress: float, total: float | None, message: str | None) -> None:
        raise NotImplementedError

    result = await session.send_request(
        types.ListRootsRequest(),
        types.ListRootsResult,
        request_read_timeout_seconds=2.5,
        progress_callback=on_progress,
    )
    assert isinstance(result, types.ListRootsResult)
    method, _params, opts = outbound.requests[0]
    assert method == "roots/list"
    assert opts == {"timeout": 2.5, "on_progress": on_progress}


@pytest.mark.anyio
async def test_send_request_omits_call_options_when_none_given():
    outbound = StubOutbound(result={"roots": []})
    session = _make_session(outbound)
    await session.send_request(types.ListRootsRequest(), types.ListRootsResult)
    _method, _params, opts = outbound.requests[0]
    assert opts is None


@pytest.mark.anyio
async def test_send_request_timeout_zero_is_forwarded():
    """0 is a real timeout (fail at the first checkpoint, `anyio.fail_after(0)`
    semantics) and must reach the channel; only `None` means "no timeout"."""
    outbound = StubOutbound(result={})
    session = _make_session(outbound)
    await session.send_request(types.PingRequest(), types.EmptyResult, request_read_timeout_seconds=0.0)
    assert outbound.requests[0][2] == {"timeout": 0.0}


@pytest.mark.anyio
async def test_send_request_without_related_id_routes_to_standalone_channel():
    """SDK-defined: no `related_request_id` routes the request onto the connection's standalone channel."""
    request_ch = StubOutbound()
    standalone_ch = StubOutbound(result={"roots": []})
    session = _two_channel_session(request_ch, standalone_ch)
    await session.send_request(types.ListRootsRequest(), types.ListRootsResult)
    assert request_ch.requests == []
    assert standalone_ch.requests[0][0] == "roots/list"


@pytest.mark.anyio
async def test_send_request_with_related_id_routes_to_request_channel():
    """SDK-defined: with `related_request_id` the request rides the per-request channel
    (the originating POST's response stream over streamable HTTP)."""
    request_ch = StubOutbound(result={"action": "cancel"})
    standalone_ch = StubOutbound()
    session = _two_channel_session(request_ch, standalone_ch)
    result = await session.send_request(
        types.ElicitRequest(params=types.ElicitRequestFormParams(message="q", requested_schema={})),
        types.ElicitResult,
        metadata=ServerMessageMetadata(related_request_id=7),
    )
    assert isinstance(result, types.ElicitResult)
    assert standalone_ch.requests == []
    assert request_ch.requests[0][0] == "elicitation/create"


@pytest.mark.anyio
async def test_send_notification_routes_by_related_request_id():
    """SDK-defined: notifications select channel by `related_request_id` exactly like requests."""
    request_ch = StubOutbound()
    standalone_ch = StubOutbound()
    session = _two_channel_session(request_ch, standalone_ch)
    await session.send_tool_list_changed()
    await session.send_progress_notification("tok", 0.5, related_request_id="req-1")
    assert [m for m, _ in standalone_ch.notifications] == ["notifications/tools/list_changed"]
    assert [m for m, _ in request_ch.notifications] == ["notifications/progress"]


@pytest.mark.anyio
async def test_report_progress_delegates_to_the_request_dispatch_context():
    """`report_progress` calls the per-request `DispatchContext.progress` seam, never the
    standalone channel: token gating and routing live in the dispatcher, not here."""
    request_ch = StubOutbound()
    standalone_ch = StubOutbound()
    session = _two_channel_session(request_ch, standalone_ch)
    await session.report_progress(0.5, total=1.0, message="halfway")
    assert request_ch.progress_calls == [(0.5, 1.0, "halfway")]
    assert standalone_ch.progress_calls == []
    assert request_ch.notifications == []


@pytest.mark.anyio
async def test_send_request_validates_the_client_result_against_the_surface_schema():
    """A spec-method result that fails the per-version surface schema raises
    `ValidationError` even when the caller's `result_type` would accept it."""
    session = _make_session(StubOutbound(result={"roots": "nope"}))
    with pytest.raises(ValidationError):
        await session.send_request(types.ListRootsRequest(), types.EmptyResult)


@pytest.mark.anyio
async def test_send_request_passes_a_spec_valid_client_result():
    """A spec-valid client result passes the surface gate and parses to the typed model."""
    session = _make_session(StubOutbound(result={"roots": [{"uri": "file:///ws"}]}))
    result = await session.send_request(types.ListRootsRequest(), types.ListRootsResult)
    assert isinstance(result, types.ListRootsResult)
    assert str(result.roots[0].uri) == "file:///ws"


@pytest.mark.anyio
async def test_send_request_skips_the_surface_gate_when_method_absent_at_version():
    """Surface row absent for the connection's version: gate is bypassed and only
    `result_type` validates."""
    session = _make_session(StubOutbound(result={}), protocol_version=LATEST_MODERN_VERSION)
    result = await session.send_request(types.PingRequest(), types.EmptyResult)
    assert isinstance(result, types.EmptyResult)


@pytest.mark.anyio
async def test_send_request_validates_result_alias_only():
    """Peer results validate alias-only; a snake_case key from the wire is
    ignored as extra, not populated by Python field name."""
    snake = {"role": "assistant", "content": {"type": "text", "text": "x"}, "model": "m", "stop_reason": "endTurn"}
    session = _make_session(StubOutbound(result=snake))
    request = types.CreateMessageRequest(params=types.CreateMessageRequestParams(messages=[], max_tokens=1))
    result = await session.send_request(request, types.CreateMessageResult)
    assert result.stop_reason is None


@pytest.mark.anyio
async def test_create_message_with_tools_returns_with_tools_result():
    outbound = StubOutbound(result={"role": "assistant", "content": [{"type": "text", "text": "ok"}], "model": "m"})
    session = _make_session(
        outbound, capabilities=ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability()))
    )
    result = await session.create_message(  # pyright: ignore[reportDeprecated]
        messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hi"))],
        max_tokens=10,
        tools=[types.Tool(name="t", input_schema={"type": "object"})],
    )
    assert isinstance(result, types.CreateMessageResultWithTools)
    method, params, _opts = outbound.requests[0]
    assert method == "sampling/createMessage"
    assert params is not None and params["tools"][0]["name"] == "t"


def test_check_client_capability_delegates_to_connection():
    outbound = StubOutbound()
    session = _make_session(outbound, capabilities=ClientCapabilities(sampling=SamplingCapability()))
    assert session.check_client_capability(ClientCapabilities(sampling=SamplingCapability())) is True
    assert session.check_client_capability(ClientCapabilities(experimental={"x": {}})) is False


def test_protocol_version_proxies_connection():
    """SDK-defined: `session.protocol_version` reads through to the held `Connection`."""
    _ARBITRARY_VERSION = "sentinel-version"  # identity-only: any string the connection holds
    conn = Connection.from_envelope(_ARBITRARY_VERSION, None, None)
    session = ServerSession(StubOutbound(), conn)
    assert session.protocol_version == _ARBITRARY_VERSION
    assert session.client_params is None