"""Tests for `OpenTelemetryMiddleware` (the context-tier OTel span middleware).

Every `Server` ships `OpenTelemetryMiddleware` at the head of `Server.middleware`,
so these tests assert against the default-configured server rather than appending
the middleware by hand.
"""

from collections.abc import Callable
from dataclasses import replace
from typing import Any

import anyio
import pytest
from mcp_types import (
    INTERNAL_ERROR,
    INVALID_PARAMS,
    CallToolRequestParams,
    CallToolResult,
    GetPromptRequestParams,
    GetPromptResult,
    ListToolsResult,
    NotificationParams,
    PaginatedRequestParams,
    RequestParamsMeta,
    Tool,
)
from opentelemetry.trace import SpanKind, StatusCode

from mcp.server._otel import OpenTelemetryMiddleware
from mcp.server.context import CallNext
from mcp.server.lowlevel.server import Server
from mcp.shared._otel import inject_trace_context, otel_span
from mcp.shared.exceptions import MCPError

from .conftest import SpanCapture
from .test_runner import Ctx, SrvT, connected_runner


@pytest.fixture
def server() -> SrvT:
    async def list_tools(ctx: Ctx, params: PaginatedRequestParams | None) -> ListToolsResult:
        return ListToolsResult(tools=[Tool(name="t", input_schema={"type": "object"})])

    return Server(name="test-server", version="0.0.1", on_list_tools=list_tools)


async def _ok_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
    return {"content": [], "isError": False}


def test_server_ships_opentelemetry_middleware_by_default() -> None:
    server = Server(name="test-server", version="0.0.1")
    assert any(isinstance(m, OpenTelemetryMiddleware) for m in server.middleware)


@pytest.mark.anyio
async def test_emits_server_span_with_method_and_target(server: SrvT, spans: SpanCapture):
    server.add_request_handler("tools/call", CallToolRequestParams, _ok_tool)
    async with connected_runner(server) as (client, _):
        spans.clear()
        result = await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
    assert result == {"content": [], "isError": False}
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.name == "tools/call mytool"
    assert span.attributes is not None
    assert span.attributes["mcp.method.name"] == "tools/call"
    assert span.attributes["gen_ai.operation.name"] == "execute_tool"
    assert span.attributes["gen_ai.tool.name"] == "mytool"
    assert isinstance(span.attributes["jsonrpc.request.id"], str)
    assert span.status.status_code == StatusCode.UNSET


@pytest.mark.anyio
async def test_tool_error_dict_result_sets_error_type(server: SrvT, spans: SpanCapture):
    async def err_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
        return {"content": [], "isError": True}

    server.add_request_handler("tools/call", CallToolRequestParams, err_tool)
    async with connected_runner(server) as (client, _):
        spans.clear()
        await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.attributes is not None
    assert span.attributes["error.type"] == "tool_error"
    assert span.status.status_code == StatusCode.ERROR


@pytest.mark.anyio
async def test_tool_error_model_result_sets_error_type(server: SrvT, spans: SpanCapture):
    async def err_tool(ctx: Ctx, params: CallToolRequestParams) -> CallToolResult:
        return CallToolResult(content=[], is_error=True)

    server.add_request_handler("tools/call", CallToolRequestParams, err_tool)
    async with connected_runner(server) as (client, _):
        spans.clear()
        await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.attributes is not None
    assert span.attributes["error.type"] == "tool_error"
    assert span.status.status_code == StatusCode.ERROR


@pytest.mark.anyio
async def test_snake_case_dict_result_is_not_a_tool_error(server: SrvT, spans: SpanCapture):
    # `is_error` is alias-only on the wire, so serialization drops it; the result reaches the
    # client as a success and the span must not contradict that.
    async def err_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
        return {"content": [], "is_error": True}

    server.add_request_handler("tools/call", CallToolRequestParams, err_tool)
    async with connected_runner(server) as (client, _):
        spans.clear()
        result = await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
    assert result == {"content": []}
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.attributes is not None
    assert "error.type" not in span.attributes
    assert span.status.status_code == StatusCode.UNSET


@pytest.mark.anyio
async def test_named_non_tool_prompt_method_omits_gen_ai_attrs(server: SrvT, spans: SpanCapture):
    async def custom(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
        return {"content": [], "isError": False}

    server.add_request_handler("custom/op", CallToolRequestParams, custom)
    async with connected_runner(server) as (client, _):
        spans.clear()
        await client.send_raw_request("custom/op", {"name": "thing", "arguments": {}})
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.name == "custom/op thing"
    assert span.attributes is not None
    assert "gen_ai.operation.name" not in span.attributes
    assert "gen_ai.tool.name" not in span.attributes
    assert "gen_ai.prompt.name" not in span.attributes


@pytest.mark.anyio
async def test_prompt_get_sets_prompt_name(server: SrvT, spans: SpanCapture):
    async def get_prompt(ctx: Ctx, params: GetPromptRequestParams) -> GetPromptResult:
        return GetPromptResult(messages=[])

    server.add_request_handler("prompts/get", GetPromptRequestParams, get_prompt)
    async with connected_runner(server) as (client, _):
        spans.clear()
        await client.send_raw_request("prompts/get", {"name": "myprompt"})
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.name == "prompts/get myprompt"
    assert span.attributes is not None
    assert span.attributes["gen_ai.prompt.name"] == "myprompt"
    assert "gen_ai.operation.name" not in span.attributes


@pytest.mark.anyio
async def test_notification_span_omits_request_id(server: SrvT, spans: SpanCapture):
    async def on_roots(ctx: Ctx, params: NotificationParams | None) -> None:
        return None

    server.add_notification_handler("notifications/roots/list_changed", NotificationParams, on_roots)
    async with connected_runner(server) as (client, _):
        spans.clear()
        await client.notify("notifications/roots/list_changed", None)
        await anyio.wait_all_tasks_blocked()
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.name == "notifications/roots/list_changed"
    assert span.attributes is not None
    assert span.attributes["mcp.method.name"] == "notifications/roots/list_changed"
    assert "jsonrpc.request.id" not in span.attributes


def _ambient(rewrite_meta: Callable[[RequestParamsMeta | None], RequestParamsMeta | None]) -> Any:
    """A middleware placed outside `OpenTelemetryMiddleware` (head of
    `Server.middleware`) that opens an ambient SERVER span around the request
    and rewrites `ctx.meta` via `rewrite_meta`, so the context-tier span sees
    the no-traceparent path yet has a current span to nest under."""

    async def middleware(ctx: Ctx, call_next: CallNext) -> Any:
        with otel_span("ambient", kind=SpanKind.SERVER):
            return await call_next(replace(ctx, meta=rewrite_meta(ctx.meta)))

    return middleware


@pytest.mark.anyio
async def test_nests_under_ambient_span_when_no_traceparent(server: SrvT, spans: SpanCapture):
    """With no `_meta` on the inbound message (a non-SDK client), the span must
    parent to the ambient current span rather than become an orphan root.
    SDK-defined: SEP-414 only covers the traceparent-present case."""

    # The in-process client always injects `_meta.traceparent`; drop it so the
    # span sees the no-carrier path.
    server.middleware.insert(0, _ambient(lambda _meta: None))
    async with connected_runner(server) as (client, _):
        spans.clear()
        await client.send_raw_request("tools/list", None)
    server_spans = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert len(server_spans) == 2
    [outer] = [s for s in server_spans if s.parent is None]
    [inner] = [s for s in server_spans if s.parent is not None]
    assert inner.context is not None and outer.context is not None
    assert inner.parent is not None
    assert inner.context.trace_id == outer.context.trace_id
    assert inner.parent.span_id == outer.context.span_id


@pytest.mark.anyio
async def test_nests_under_ambient_span_when_meta_lacks_traceparent(server: SrvT, spans: SpanCapture):
    """`_meta` is present but carries no `traceparent` (e.g. only a
    `progressToken`). `extract()` would yield an empty Context here, which
    would orphan the span; the middleware must fall through to ambient
    parenting just as if `_meta` were absent."""

    server.middleware.insert(0, _ambient(lambda _meta: {"progressToken": "tok"}))
    async with connected_runner(server) as (client, _):
        spans.clear()
        await client.send_raw_request("tools/list", None)
    server_spans = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert len(server_spans) == 2
    [outer] = [s for s in server_spans if s.parent is None]
    [inner] = [s for s in server_spans if s.parent is not None]
    assert inner.context is not None and outer.context is not None
    assert inner.parent is not None
    assert inner.context.trace_id == outer.context.trace_id
    assert inner.parent.span_id == outer.context.span_id


@pytest.mark.anyio
async def test_extracts_trace_context_from_meta(server: SrvT, spans: SpanCapture):
    meta: dict[str, Any] = {}
    inject_trace_context(meta)
    async with connected_runner(server) as (client, _):
        spans.clear()
        await client.send_raw_request("tools/list", {"_meta": meta})
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.parent is not None


@pytest.mark.anyio
async def test_records_error_status_on_mcp_error(server: SrvT, spans: SpanCapture):
    async with connected_runner(server) as (client, _):
        spans.clear()
        with pytest.raises(MCPError) as exc:
            await client.send_raw_request("resources/list", None)
        assert exc.value.error.code != 0
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.status.status_code == StatusCode.ERROR
    assert span.status.description == "Method not found"
    assert span.attributes is not None
    assert span.attributes["error.type"] == str(exc.value.error.code)
    assert span.attributes["rpc.response.status_code"] == str(exc.value.error.code)
    assert not [e for e in span.events if e.name == "exception"]


@pytest.mark.anyio
async def test_validation_failure_sets_sanitized_status(server: SrvT, spans: SpanCapture):
    server.add_request_handler("tools/call", CallToolRequestParams, _ok_tool)
    async with connected_runner(server) as (client, _):
        spans.clear()
        with pytest.raises(MCPError):
            await client.send_raw_request("tools/call", {"name": 123})
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.status.status_code == StatusCode.ERROR
    assert span.status.description == "Invalid request parameters"
    assert span.attributes is not None
    assert span.attributes["error.type"] == str(INVALID_PARAMS)
    assert span.attributes["rpc.response.status_code"] == str(INVALID_PARAMS)
    assert span.attributes["gen_ai.operation.name"] == "execute_tool"
    assert "gen_ai.tool.name" not in span.attributes
    assert not span.events


@pytest.mark.anyio
async def test_records_error_status_on_handler_exception(server: SrvT, spans: SpanCapture):
    async def failing(ctx: Ctx, params: PaginatedRequestParams | None) -> Any:
        raise ValueError("handler blew up")

    server.add_request_handler("tools/list", PaginatedRequestParams, failing)
    async with connected_runner(server) as (client, _):
        spans.clear()
        with pytest.raises(MCPError):
            await client.send_raw_request("tools/list", None)
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.status.status_code == StatusCode.ERROR
    assert span.status.description == "handler blew up"
    assert span.attributes is not None
    assert span.attributes["error.type"] == "ValueError"
    [event] = [e for e in span.events if e.name == "exception"]
    assert event.attributes is not None
    assert event.attributes["exception.type"] == "ValueError"


@pytest.mark.anyio
async def test_records_error_status_on_malformed_spec_result(server: SrvT, spans: SpanCapture):
    """Result serialization runs inside the span, so a handler returning a
    malformed dict for a spec method (INTERNAL_ERROR on the wire) is recorded
    on the span rather than closing it as a success."""

    async def bad_result(ctx: Ctx, params: PaginatedRequestParams | None) -> dict[str, Any]:
        return {"tools": 42}

    server.add_request_handler("tools/list", PaginatedRequestParams, bad_result)
    async with connected_runner(server) as (client, _):
        spans.clear()
        with pytest.raises(MCPError) as exc:
            await client.send_raw_request("tools/list", None)
        assert exc.value.error.code == INTERNAL_ERROR
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.status.status_code == StatusCode.ERROR
    assert span.attributes is not None
    assert span.attributes["error.type"] == str(INTERNAL_ERROR)
    assert span.attributes["rpc.response.status_code"] == str(INTERNAL_ERROR)


@pytest.mark.anyio
async def test_passes_rewritten_context_through(server: SrvT, spans: SpanCapture):
    seen_arguments: dict[str, Any] = {}

    async def call_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
        seen_arguments.update(params.arguments or {})
        return {"content": [], "isError": False}

    async def inject_arg(ctx: Ctx, call_next: CallNext) -> Any:
        assert ctx.params is not None
        arguments = {**ctx.params.get("arguments", {}), "injected": True}
        return await call_next(replace(ctx, params={**ctx.params, "arguments": arguments}))

    server.add_request_handler("tools/call", CallToolRequestParams, call_tool)
    server.middleware.append(inject_arg)
    async with connected_runner(server) as (client, _):
        spans.clear()
        await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {"x": 1}})
    assert seen_arguments == {"x": 1, "injected": True}
    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
    assert span.name == "tools/call mytool"
