from __future__ import annotations

import asyncio
import contextlib
import socket
from typing import TYPE_CHECKING, Any

import pytest
import uvicorn
from fastmcp import Client, FastMCP
from fastmcp.client.transports import StreamableHttpTransport
from fastmcp.server.dependencies import get_http_headers

from beaver_gateway.backends.claude_sdk import ORIGIN_HEADER
from beaver_gateway.mcp.internal_app import build_internal_app
from beaver_gateway.mcp.types import McpServer

if TYPE_CHECKING:
    from collections.abc import AsyncIterator


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


@contextlib.asynccontextmanager
async def _serving(app: Any, port: int) -> AsyncIterator[None]:
    server = uvicorn.Server(
        uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning")
    )
    task = asyncio.create_task(server.serve())
    while not server.started:
        await asyncio.sleep(0.02)
    try:
        yield
    finally:
        server.should_exit = True
        await task


@pytest.mark.anyio
async def test_the_origin_header_reaches_the_upstream_mcp() -> None:
    upstream = FastMCP("upstream")

    @upstream.tool
    def caller() -> str:
        return get_http_headers().get(ORIGIN_HEADER.lower(), "")

    up, gateway = _free_port(), _free_port()
    app, urls, _ = build_internal_app(
        [McpServer.http(name="t3", url=f"http://127.0.0.1:{up}/mcp")],
        host="127.0.0.1",
        port=gateway,
    )
    seen = []
    async with (
        _serving(upstream.http_app(transport="http", path="/mcp"), up),
        _serving(app, gateway),
    ):
        for conversation in ("conv-a", "conv-b"):
            transport = StreamableHttpTransport(
                urls["t3"], headers={ORIGIN_HEADER: conversation}
            )
            async with Client(transport) as client:
                result = await client.call_tool("caller", {})
                seen.append(result.data)
    assert seen == ["conv-a", "conv-b"]