"""Tests for `ClientPeer`.

Each typed method is tested by wrapping a `DirectDispatcher` in `ClientPeer`,
calling it, and asserting (a) the right method+params went out and (b) the
return value is the typed result model.
"""

from collections.abc import Mapping
from typing import Any

import anyio
import pytest

from mcp.shared.dispatcher import DispatchContext
from mcp.shared.peer import ClientPeer, dump_params
from mcp.shared.transport_context import TransportContext
from mcp.types import (
    CreateMessageResult,
    CreateMessageResultWithTools,
    ElicitResult,
    ListRootsResult,
    SamplingMessage,
    TextContent,
    Tool,
)

from .conftest import direct_pair
from .test_dispatcher import running_pair

DCtx = DispatchContext[TransportContext]


class _Recorder:
    def __init__(self, result: dict[str, Any]) -> None:
        self.result = result
        self.seen: list[tuple[str, Mapping[str, Any] | None]] = []

    async def on_request(self, ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
        self.seen.append((method, params))
        return self.result


@pytest.mark.anyio
async def test_peer_sample_sends_create_message_and_returns_typed_result():
    rec = _Recorder({"role": "assistant", "content": {"type": "text", "text": "hi"}, "model": "m"})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.sample(
                [SamplingMessage(role="user", content=TextContent(type="text", text="hello"))],
                max_tokens=10,
            )
        method, params = rec.seen[0]
        assert method == "sampling/createMessage"
        assert params is not None and params["maxTokens"] == 10
        assert isinstance(result, CreateMessageResult)
        assert result.model == "m"


@pytest.mark.anyio
async def test_peer_sample_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"}
    rec = _Recorder(snake)
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.sample(
                [SamplingMessage(role="user", content=TextContent(type="text", text="q"))], max_tokens=1
            )
        assert isinstance(result, CreateMessageResult)
        assert result.stop_reason is None


@pytest.mark.anyio
async def test_peer_sample_with_tools_returns_with_tools_result():
    rec = _Recorder({"role": "assistant", "content": [{"type": "text", "text": "x"}], "model": "m"})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.sample(
                [SamplingMessage(role="user", content=TextContent(type="text", text="q"))],
                max_tokens=5,
                tools=[Tool(name="t", input_schema={"type": "object"})],
            )
        method, params = rec.seen[0]
        assert method == "sampling/createMessage"
        assert params is not None and params["tools"][0]["name"] == "t"
        assert isinstance(result, CreateMessageResultWithTools)


@pytest.mark.anyio
async def test_peer_elicit_form_sends_elicitation_create_with_form_params():
    rec = _Recorder({"action": "accept", "content": {"name": "Max"}})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.elicit_form("Your name?", requested_schema={"type": "object", "properties": {}})
        method, params = rec.seen[0]
        assert method == "elicitation/create"
        assert params is not None and params["mode"] == "form"
        assert params["message"] == "Your name?"
        assert isinstance(result, ElicitResult)


@pytest.mark.anyio
async def test_peer_elicit_url_sends_elicitation_create_with_url_params():
    rec = _Recorder({"action": "accept"})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.elicit_url("Auth needed", url="https://example.com/auth", elicitation_id="e1")
        method, params = rec.seen[0]
        assert method == "elicitation/create"
        assert params is not None and params["mode"] == "url"
        assert params["url"] == "https://example.com/auth"
        assert isinstance(result, ElicitResult)


@pytest.mark.anyio
async def test_peer_list_roots_sends_roots_list_and_returns_typed_result():
    rec = _Recorder({"roots": [{"uri": "file:///workspace"}]})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.list_roots()
        method, _ = rec.seen[0]
        assert method == "roots/list"
        assert isinstance(result, ListRootsResult)
        assert len(result.roots) == 1
        assert str(result.roots[0].uri) == "file:///workspace"


@pytest.mark.anyio
async def test_peer_list_roots_with_meta_sends_meta_in_params():
    rec = _Recorder({"roots": []})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            await peer.list_roots(meta={"traceId": "t1"})
        method, params = rec.seen[0]
        assert method == "roots/list"
        assert params == {"_meta": {"traceId": "t1"}}


def test_dump_params_merges_meta_over_model_meta():
    out = dump_params(None, None)
    assert out is None
    out = dump_params(None, {"k": 1})
    assert out == {"_meta": {"k": 1}}


def test_dump_params_serializes_meta_by_alias():
    """`progress_token` (the Python key an inbound `ctx.meta` carries) emits
    its wire alias `progressToken`; undeclared keys pass through unchanged."""
    out = dump_params(None, {"progress_token": 7, "traceparent": "00-abc"})
    assert out == {"_meta": {"progressToken": 7, "traceparent": "00-abc"}}
    # The wire spelling is already canonical and survives as-is.
    out = dump_params(None, {"progressToken": "tok"})
    assert out == {"_meta": {"progressToken": "tok"}}


@pytest.mark.anyio
async def test_peer_notify_forwards_to_wrapped_outbound():
    sent: list[tuple[str, Mapping[str, Any] | None]] = []

    class _Out:
        async def send_raw_request(
            self, method: str, params: Mapping[str, Any] | None, opts: Any = None
        ) -> dict[str, Any]:
            raise NotImplementedError

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

    await ClientPeer(_Out()).notify("n", {"x": 1})
    assert sent == [("n", {"x": 1})]


@pytest.mark.anyio
async def test_peer_ping_sends_ping_and_returns_none():
    rec = _Recorder({})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.ping()
        method, _ = rec.seen[0]
        assert method == "ping"
        assert result is None