import asyncio
import socket
import subprocess
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, AsyncGenerator

import pytest
from mcp import ClientSession
from mcp.types import CreateMessageResult, TextContent
from opentelemetry.trace import Tracer

from tests.collector import OTLPServer, Telemetry


# The way MCP SDK creates async tasks means we need this to be called inline with the test,
# not as a fixture.
@asynccontextmanager
async def mcp_client(
    transport: str, tracer: Tracer, otlp_endpoint: str
) -> AsyncGenerator[ClientSession, None]:
    # Lazy import to get instrumented versions. Users will use opentelemetry-instrument or otherwise
    # initialize instrumentation as early as possible and should not run into issues, but we control
    # instrumentation through fixtures instead.
    from mcp.client.sse import sse_client
    from mcp.client.stdio import StdioServerParameters, stdio_client
    from mcp.client.streamable_http import streamable_http_client

    async def sampling_callback(context: Any, params: Any) -> CreateMessageResult:
        del context, params
        with tracer.start_as_current_span("whoami"):
            return CreateMessageResult(
                role="assistant",
                content=TextContent(type="text", text="OpenInference"),
                model="test-model",
            )

    server_script = str(Path(__file__).parent / "mcpserver.py")
    pythonpath = str(Path(__file__).parent.parent)

    def _choose_free_port() -> int:
        with socket.socket() as s:
            s.bind(("127.0.0.1", 0))
            return s.getsockname()[1]  # type: ignore

    async def _wait_for_port(
        host: str, port: int, attempts: int = 100, delay: float = 0.05
    ) -> None:
        for _ in range(attempts):
            try:
                reader, writer = await asyncio.open_connection(host, port)
                writer.close()
                await writer.wait_closed()
                return
            except Exception:
                await asyncio.sleep(delay)
        raise RuntimeError(f"Timed out waiting for server on {host}:{port}")

    match transport:
        case "stdio":
            async with stdio_client(
                StdioServerParameters(
                    command=sys.executable,
                    args=[server_script],
                    env={
                        "MCP_TRANSPORT": "stdio",
                        "OTEL_EXPORTER_OTLP_ENDPOINT": otlp_endpoint,
                        "PYTHONPATH": pythonpath,
                    },
                )
            ) as (reader, writer), ClientSession(
                reader, writer, sampling_callback=sampling_callback
            ) as client:
                await client.initialize()
                yield client
        case "sse":
            port = _choose_free_port()
            proc = await asyncio.create_subprocess_exec(
                sys.executable,
                server_script,
                env={
                    "MCP_TRANSPORT": "sse",
                    "MCP_PORT": str(port),
                    "OTEL_EXPORTER_OTLP_ENDPOINT": otlp_endpoint,
                    "PYTHONPATH": pythonpath,
                },
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            try:
                await _wait_for_port("127.0.0.1", port)
                async with sse_client(f"http://localhost:{port}/sse") as (
                    reader,
                    writer,
                ), ClientSession(reader, writer, sampling_callback=sampling_callback) as client:
                    await client.initialize()
                    yield client
            finally:
                proc.kill()
                await proc.wait()
        case "streamable-http":
            port = _choose_free_port()
            proc = await asyncio.create_subprocess_exec(
                sys.executable,
                server_script,
                env={
                    "MCP_TRANSPORT": "streamable-http",
                    "MCP_PORT": str(port),
                    "OTEL_EXPORTER_OTLP_ENDPOINT": otlp_endpoint,
                    "PYTHONPATH": pythonpath,
                },
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            try:
                await _wait_for_port("127.0.0.1", port)
                async with streamable_http_client(f"http://localhost:{port}/mcp") as streams:
                    reader, writer = streams[:2]
                    async with ClientSession(
                        reader, writer, sampling_callback=sampling_callback
                    ) as client:
                        await client.initialize()
                        yield client
            finally:
                proc.kill()
                await proc.wait()


@pytest.mark.parametrize("transport", ["sse", "stdio", "streamable-http"])
async def test_hello(
    transport: str, tracer: Tracer, telemetry: Telemetry, otlp_collector: OTLPServer
) -> None:
    async with mcp_client(
        transport, tracer, f"http://localhost:{otlp_collector.server_port}/"
    ) as client:
        with tracer.start_as_current_span("root"):
            tools_res = await client.list_tools()
            assert len(tools_res.tools) == 1
            assert tools_res.tools[0].name == "hello"
            tool_res = await client.call_tool("hello")
            content = tool_res.content[0]
            assert isinstance(content, TextContent)
            assert content.text == "Hello OpenInference!"

    for resource_spans in telemetry.traces:
        for scope_spans in resource_spans.scope_spans:
            match scope_spans.scope.name:
                case "mcp-test-client":
                    for span in scope_spans.spans:
                        match span.name:
                            case "root":
                                root_span = span
                            case "whoami":
                                whoami_span = span
                case "mcp-test-server":
                    server_span = scope_spans.spans[0]
    assert root_span.name == "root"
    assert server_span.name == "hello"
    assert whoami_span.name == "whoami"
    assert server_span.trace_id == root_span.trace_id
    assert server_span.parent_span_id == root_span.span_id
    assert whoami_span.trace_id == root_span.trace_id
    assert whoami_span.parent_span_id == server_span.span_id


async def test_stdio_validation_error(tracer: Tracer, otlp_collector: OTLPServer) -> None:
    """Test that ValidationError in stdio transport doesn't crash the instrumentation."""
    # Lazy import to get instrumented versions. Users will use opentelemetry-instrument or otherwise
    # initialize instrumentation as early as possible and should not run into issues, but we control
    # instrumentation through fixtures instead.
    from mcp.client.stdio import StdioServerParameters, stdio_client
    from pydantic_core import ValidationError

    validation_error_received = False

    async def message_handler(message: Any) -> None:
        nonlocal validation_error_received
        if isinstance(message, ValidationError):
            validation_error_received = True

    server_script = str(Path(__file__).parent / "mcpserver_invalid.py")
    pythonpath = str(Path(__file__).parent.parent)

    async with stdio_client(
        StdioServerParameters(
            command=sys.executable,
            args=[server_script],
            env={
                "MCP_TRANSPORT": "stdio",
                "OTEL_EXPORTER_OTLP_ENDPOINT": f"http://localhost:{otlp_collector.server_port}/",
                "PYTHONPATH": pythonpath,
            },
        )
    ) as (reader, writer), ClientSession(reader, writer, message_handler=message_handler):
        # The server will send an invalid message that triggers ValidationError
        # Without the fix, this causes AttributeError: ValidationError has no attribute 'message'
        # With the fix, it checks isinstance(item, SessionMessage) first
        await asyncio.sleep(0.5)

    # Test passes if we didn't crash with AttributeError
    assert validation_error_received


async def test_stream_writer_exception_handling(tracer: Tracer) -> None:
    """Test that InstrumentedStreamWriter correctly handles exceptions passed through the stream."""
    # Lazy import to get instrumented versions
    from anyio import create_memory_object_stream
    from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
    from pydantic_core import ValidationError

    from openinference.instrumentation.mcp import InstrumentedStreamWriter

    # Create a memory stream for testing
    write_stream: MemoryObjectSendStream[Any]
    read_stream: MemoryObjectReceiveStream[Any]
    write_stream, read_stream = create_memory_object_stream(10)

    # Wrap the write stream with instrumentation
    instrumented_writer = InstrumentedStreamWriter(write_stream)

    # Create a validation error to simulate what MCP does
    validation_error = ValidationError.from_exception_data(
        "Test validation error",
        [
            {
                "type": "missing",
                "loc": ("method",),
                "input": {},
            }
        ],
    )

    # Without the fix, this would cause AttributeError when the instrumentation
    # tries to cast the exception to SessionMessage and access .message.root
    async with instrumented_writer, read_stream:
        await instrumented_writer.send(validation_error)

        # Verify the exception was passed through correctly
        received = await read_stream.receive()
        assert isinstance(received, ValidationError)
