import json
from collections.abc import AsyncIterator, Callable
from contextlib import asynccontextmanager
from typing import cast

import httpx
import pytest
from fastapi import FastAPI
from starlette.requests import ClientDisconnect
from starlette.types import Message, Scope

from app.gateway import MAX_REQUEST_BODY_BYTES, create_gateway_app

type Handler = Callable[[httpx.Request], httpx.Response]


class _TwoEventStream(httpx.AsyncByteStream):
    async def __aiter__(self) -> AsyncIterator[bytes]:
        yield b"data: first\n\n"
        yield b"data: second\n\n"


class _TrackedStream(httpx.AsyncByteStream):
    def __init__(self) -> None:
        self.closed = False
        self.started = False

    async def __aiter__(self) -> AsyncIterator[bytes]:
        self.started = True
        yield b"data: never-delivered\n\n"

    async def aclose(self) -> None:
        self.closed = True


@asynccontextmanager
async def _gateway_client(handler: Handler) -> AsyncIterator[httpx.AsyncClient]:
    gateway = create_gateway_app(httpx.MockTransport(handler))
    async with gateway.router.lifespan_context(gateway):
        transport = httpx.ASGITransport(app=gateway)
        async with httpx.AsyncClient(transport=transport, base_url="http://gateway") as client:
            yield client


@pytest.mark.asyncio
async def test_gateway_live_and_docs_are_disabled() -> None:
    gateway = create_gateway_app(httpx.MockTransport(lambda _: httpx.Response(500)))
    try:
        assert gateway.docs_url is None
        assert gateway.redoc_url is None
        assert gateway.openapi_url is None

        transport = httpx.ASGITransport(app=gateway)
        async with httpx.AsyncClient(transport=transport, base_url="http://gateway") as client:
            response = await client.get("/gateway/live")
        assert response.status_code == 200
        assert response.json() == {"status": "ok"}
    finally:
        await _close_gateway(gateway)


async def _close_gateway(gateway: FastAPI) -> None:
    async with gateway.router.lifespan_context(gateway):
        pass


@pytest.mark.asyncio
async def test_successful_get_preserves_path_query_and_allowed_header() -> None:
    seen: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen.append(request)
        return httpx.Response(200, json={"result": "ok"}, headers={"x-request-id": "req-1"})

    async with _gateway_client(handler) as client:
        response = await client.get(
            "/api/v1/search?q=granite%20deposit&limit=5",
            headers={"x-request-id": "req-1", "x-forwarded-host": "evil.invalid"},
        )

    assert response.status_code == 200
    assert response.json() == {"result": "ok"}
    assert response.headers["x-request-id"] == "req-1"
    assert seen[0].url == "http://api:8000/api/v1/search?q=granite%20deposit&limit=5"
    assert seen[0].headers["host"] == "api:8000"
    assert "x-forwarded-host" not in seen[0].headers


@pytest.mark.asyncio
async def test_post_json_body_and_content_type_are_forwarded() -> None:
    seen: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen.append(request)
        return httpx.Response(201, json={"accepted": True})

    async with _gateway_client(handler) as client:
        response = await client.post("/api/v1/query", json={"question": "Where is gold?"})

    assert response.status_code == 201
    assert response.json() == {"accepted": True}
    assert seen[0].method == "POST"
    assert seen[0].headers["content-type"] == "application/json"
    assert json.loads(seen[0].content) == {"question": "Where is gold?"}


@pytest.mark.asyncio
async def test_upstream_origin_is_fixed_against_host_and_path_ssrf_attempts() -> None:
    seen: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen.append(request)
        return httpx.Response(200)

    async with _gateway_client(handler) as client:
        response = await client.get(
            "//evil.invalid/http://metadata.internal/latest?target=http://attacker.invalid",
            headers={"host": "attacker.invalid", "forwarded": "host=attacker.invalid"},
        )

    assert response.status_code == 200
    assert seen[0].url.scheme == "http"
    assert seen[0].url.host == "api"
    assert seen[0].url.port == 8000
    assert seen[0].headers["host"] == "api:8000"
    assert "forwarded" not in seen[0].headers


@pytest.mark.asyncio
async def test_request_larger_than_one_mib_is_rejected_before_upstream() -> None:
    upstream_calls = 0

    def handler(_: httpx.Request) -> httpx.Response:
        nonlocal upstream_calls
        upstream_calls += 1
        return httpx.Response(200)

    async with _gateway_client(handler) as client:
        response = await client.post("/upload", content=b"x" * (MAX_REQUEST_BODY_BYTES + 1))

    assert response.status_code == 413
    assert response.json() == {"detail": "request body too large"}
    assert upstream_calls == 0


@pytest.mark.asyncio
async def test_upstream_transport_error_returns_redacted_502() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        raise httpx.ConnectError(
            "api:8000 refused connection with internal-token-value",
            request=request,
        )

    async with _gateway_client(handler) as client:
        response = await client.get("/api/v1/health/ready")

    assert response.status_code == 502
    assert response.json() == {"detail": "upstream unavailable"}
    assert "api:8000" not in response.text
    assert "internal-token-value" not in response.text


@pytest.mark.asyncio
async def test_sensitive_upstream_response_headers_are_not_forwarded() -> None:
    def handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            content=b"ok",
            headers={
                "content-type": "text/plain",
                "cache-control": "no-store",
                "set-cookie": "session=private",
                "server": "internal-api",
                "x-internal-token": "private",
            },
        )

    async with _gateway_client(handler) as client:
        response = await client.get("/status")

    assert response.headers["content-type"] == "text/plain"
    assert response.headers["cache-control"] == "no-store"
    assert "set-cookie" not in response.headers
    assert "server" not in response.headers
    assert "x-internal-token" not in response.headers


@pytest.mark.asyncio
async def test_sse_events_are_forwarded_as_separate_asgi_body_chunks() -> None:
    def handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            headers={"content-type": "text/event-stream"},
            stream=_TwoEventStream(),
        )

    gateway = create_gateway_app(httpx.MockTransport(handler))
    sent: list[Message] = []
    request_delivered = False

    async def receive() -> Message:
        nonlocal request_delivered
        if not request_delivered:
            request_delivered = True
            return {"type": "http.request", "body": b"", "more_body": False}
        return {"type": "http.disconnect"}

    async def send(message: Message) -> None:
        sent.append(message)

    scope = cast(
        Scope,
        {
            "type": "http",
            "asgi": {"version": "3.0", "spec_version": "2.4"},
            "http_version": "1.1",
            "method": "GET",
            "scheme": "http",
            "path": "/api/v1/chat/stream",
            "raw_path": b"/api/v1/chat/stream",
            "query_string": b"",
            "headers": [(b"accept", b"text/event-stream")],
            "client": ("127.0.0.1", 12345),
            "server": ("gateway", 80),
        },
    )

    async with gateway.router.lifespan_context(gateway):
        await gateway(scope, receive, send)

    body_messages = [message for message in sent if message["type"] == "http.response.body"]
    assert [message["body"] for message in body_messages] == [
        b"data: first\n\n",
        b"data: second\n\n",
        b"",
    ]
    assert [message["more_body"] for message in body_messages] == [True, True, False]


@pytest.mark.asyncio
async def test_upstream_stream_closes_when_downstream_disconnects_before_body_iteration() -> None:
    upstream_stream = _TrackedStream()

    def handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            headers={"content-type": "text/event-stream"},
            stream=upstream_stream,
        )

    gateway = create_gateway_app(httpx.MockTransport(handler))
    request_delivered = False

    async def receive() -> Message:
        nonlocal request_delivered
        if not request_delivered:
            request_delivered = True
            return {"type": "http.request", "body": b"", "more_body": False}
        return {"type": "http.disconnect"}

    async def disconnect_on_response_start(message: Message) -> None:
        if message["type"] == "http.response.start":
            raise OSError("downstream disconnected")

    scope = cast(
        Scope,
        {
            "type": "http",
            "asgi": {"version": "3.0", "spec_version": "2.4"},
            "http_version": "1.1",
            "method": "GET",
            "scheme": "http",
            "path": "/api/v1/chat/stream",
            "raw_path": b"/api/v1/chat/stream",
            "query_string": b"",
            "headers": [(b"accept", b"text/event-stream")],
            "client": ("127.0.0.1", 12345),
            "server": ("gateway", 80),
        },
    )

    async with gateway.router.lifespan_context(gateway):
        with pytest.raises(ClientDisconnect):
            await gateway(scope, receive, disconnect_on_response_start)

    assert upstream_stream.closed is True
    assert upstream_stream.started is False