import json

import pytest
from fastapi import FastAPI, Request
from mcp_types import TextResourceContents

from fastmcp import Client, FastMCP
from fastmcp.client.transports import SSETransport, StreamableHttpTransport
from fastmcp.server import create_proxy
from fastmcp.server.providers.openapi import MCPType, RouteMap
from fastmcp.utilities.tests import run_server_async


def create_fastmcp_server_for_headers() -> FastMCP:
    """Create a FastMCP server from FastAPI app with experimental parser."""
    app = FastAPI()

    @app.get("/headers")
    def get_headers(request: Request):
        return request.headers

    @app.get("/headers/{header_name}")
    def get_header_by_name(header_name: str, request: Request):
        return request.headers[header_name]

    @app.post("/headers")
    def post_headers(request: Request):
        return request.headers

    mcp = FastMCP.from_fastapi(
        app,
        httpx_client_kwargs={"headers": {"x-server-header": "test-abc"}},
        route_maps=[
            # GET requests with path parameters go to ResourceTemplate
            RouteMap(
                methods=["GET"],
                pattern=r".*\{.*\}.*",
                mcp_type=MCPType.RESOURCE_TEMPLATE,
            ),
            # GET requests without path parameters go to Resource
            RouteMap(methods=["GET"], pattern=r".*", mcp_type=MCPType.RESOURCE),
        ],
    )

    return mcp


@pytest.fixture
async def shttp_server():
    """Start a test server with StreamableHttp transport."""
    server = create_fastmcp_server_for_headers()
    async with run_server_async(server, transport="http") as url:
        yield url


@pytest.fixture
async def sse_server():
    """Start a test server with SSE transport."""
    server = create_fastmcp_server_for_headers()
    async with run_server_async(server, transport="sse") as url:
        yield url


@pytest.fixture
async def proxy_server(shttp_server: str):
    """Start a proxy server."""
    proxy = create_proxy(StreamableHttpTransport(shttp_server))
    async with run_server_async(proxy, transport="http") as url:
        yield url


async def test_fastapi_client_headers_streamable_http_resource(shttp_server: str):
    async with Client(transport=StreamableHttpTransport(shttp_server)) as client:
        result = await client.read_resource("resource://get_headers_headers_get")
        assert isinstance(result[0], TextResourceContents)
        headers = json.loads(result[0].text)
        assert headers["x-server-header"] == "test-abc"


async def test_fastapi_client_headers_sse_resource(sse_server: str):
    async with Client(transport=SSETransport(sse_server)) as client:
        result = await client.read_resource("resource://get_headers_headers_get")
        assert isinstance(result[0], TextResourceContents)
        headers = json.loads(result[0].text)
        assert headers["x-server-header"] == "test-abc"


async def test_fastapi_client_headers_streamable_http_tool(shttp_server: str):
    async with Client(transport=StreamableHttpTransport(shttp_server)) as client:
        result = await client.call_tool("post_headers_headers_post")
        headers: dict[str, str] = result.data
        assert headers["x-server-header"] == "test-abc"


async def test_fastapi_client_headers_sse_tool(sse_server: str):
    async with Client(transport=SSETransport(sse_server)) as client:
        result = await client.call_tool("post_headers_headers_post")
        headers: dict[str, str] = result.data
        assert headers["x-server-header"] == "test-abc"


async def test_client_headers_sse_resource(sse_server: str):
    async with Client(
        transport=SSETransport(sse_server, headers={"X-TEST": "test-123"})
    ) as client:
        result = await client.read_resource("resource://get_headers_headers_get")
        assert isinstance(result[0], TextResourceContents)
        headers = json.loads(result[0].text)
        assert headers["x-test"] == "test-123"


async def test_client_headers_shttp_resource(shttp_server: str):
    async with Client(
        transport=StreamableHttpTransport(shttp_server, headers={"X-TEST": "test-123"})
    ) as client:
        result = await client.read_resource("resource://get_headers_headers_get")
        assert isinstance(result[0], TextResourceContents)
        headers = json.loads(result[0].text)
        assert headers["x-test"] == "test-123"


async def test_client_headers_sse_resource_template(sse_server: str):
    async with Client(
        transport=SSETransport(sse_server, headers={"X-TEST": "test-123"})
    ) as client:
        result = await client.read_resource(
            "resource://get_header_by_name_headers/x-test"
        )
        assert isinstance(result[0], TextResourceContents)
        header = json.loads(result[0].text)
        assert header == "test-123"


async def test_client_headers_shttp_resource_template(shttp_server: str):
    async with Client(
        transport=StreamableHttpTransport(shttp_server, headers={"X-TEST": "test-123"})
    ) as client:
        result = await client.read_resource(
            "resource://get_header_by_name_headers/x-test"
        )
        assert isinstance(result[0], TextResourceContents)
        header = json.loads(result[0].text)
        assert header == "test-123"


async def test_client_headers_sse_tool(sse_server: str):
    async with Client(
        transport=SSETransport(sse_server, headers={"X-TEST": "test-123"})
    ) as client:
        result = await client.call_tool("post_headers_headers_post")
        headers: dict[str, str] = result.data
        assert headers["x-test"] == "test-123"


async def test_client_headers_shttp_tool(shttp_server: str):
    async with Client(
        transport=StreamableHttpTransport(shttp_server, headers={"X-TEST": "test-123"})
    ) as client:
        result = await client.call_tool("post_headers_headers_post")
        headers: dict[str, str] = result.data
        assert headers["x-test"] == "test-123"


async def test_client_overrides_server_headers(shttp_server: str):
    async with Client(
        transport=StreamableHttpTransport(
            shttp_server, headers={"x-server-header": "test-client"}
        )
    ) as client:
        result = await client.read_resource("resource://get_headers_headers_get")
        assert isinstance(result[0], TextResourceContents)
        headers = json.loads(result[0].text)
        assert headers["x-server-header"] == "test-client"


async def test_client_with_excluded_header_is_ignored(sse_server: str):
    async with Client(
        transport=SSETransport(
            sse_server,
            headers={
                "x-server-header": "test-client",
                "host": "1.2.3.4",
                "not-host": "1.2.3.4",
            },
        )
    ) as client:
        result = await client.read_resource("resource://get_headers_headers_get")
        assert isinstance(result[0], TextResourceContents)
        headers = json.loads(result[0].text)
        assert headers["not-host"] == "1.2.3.4"
        assert headers["host"] == "fastapi"


@pytest.mark.flaky(retries=2, delay=1)
async def test_client_headers_proxy(proxy_server: str):
    """
    Test that client headers are passed through the proxy to the remove server.
    """
    async with Client(transport=StreamableHttpTransport(proxy_server)) as client:
        result = await client.read_resource("resource://get_headers_headers_get")
        assert isinstance(result[0], TextResourceContents)
        headers = json.loads(result[0].text)
        assert headers["x-server-header"] == "test-abc"