from __future__ import annotations

import asyncio
import socket

import pytest
import uvicorn
from mcp import Client
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamable_http_client
from mcp.server.mcpserver import MCPServer
from starlette.applications import Starlette


def _free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("127.0.0.1", 0))
        return int(sock.getsockname()[1])


def _dual_transport_app(server: MCPServer) -> Starlette:
    modern = server.streamable_http_app(host="127.0.0.1")
    legacy = server.sse_app(host="127.0.0.1")
    return Starlette(
        routes=[*modern.routes, *legacy.routes],
        lifespan=modern.router.lifespan_context,
    )


async def _wait_for_server(port: int) -> None:
    for _ in range(100):
        try:
            reader, writer = await asyncio.open_connection("127.0.0.1", port)
        except OSError:
            await asyncio.sleep(0.02)
            continue
        writer.close()
        await writer.wait_closed()
        return
    raise RuntimeError("local test server did not start")


async def _exercise_fresh_dual_transport_server() -> None:
    server = MCPServer("dual-transport-proof")

    @server.tool()
    async def identity() -> str:
        return "dual-ok"

    app = _dual_transport_app(server)
    port = _free_port()
    uvicorn_server = uvicorn.Server(
        uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning")
    )
    task = asyncio.create_task(uvicorn_server.serve())
    await _wait_for_server(port)

    try:
        async with Client(
            streamable_http_client(f"http://127.0.0.1:{port}/mcp"),
            mode="auto",
            cache=None,
        ) as modern_client:
            modern_tools = await modern_client.list_tools()
            modern_result = await modern_client.call_tool("identity", {})
            assert modern_client.protocol_version == "2026-07-28"
            assert any(tool.name == "identity" for tool in modern_tools.tools)
            assert modern_result.content[0].text == "dual-ok"

        async with Client(
            sse_client(f"http://127.0.0.1:{port}/sse"),
            mode="legacy",
            cache=None,
        ) as legacy_client:
            legacy_tools = await legacy_client.list_tools()
            legacy_result = await legacy_client.call_tool("identity", {})
            assert any(tool.name == "identity" for tool in legacy_tools.tools)
            assert legacy_result.content[0].text == "dual-ok"

        async with Client(
            streamable_http_client(f"http://127.0.0.1:{port}/mcp"),
            mode="auto",
            cache=None,
        ) as second_modern_client:
            second_result = await second_modern_client.call_tool("identity", {})
            assert second_modern_client.protocol_version == "2026-07-28"
            assert second_result.content[0].text == "dual-ok"
    finally:
        uvicorn_server.should_exit = True
        await task


@pytest.mark.asyncio
async def test_one_server_instance_supports_modern_and_legacy_lifecycles() -> None:
    await _exercise_fresh_dual_transport_server()
    await _exercise_fresh_dual_transport_server()