"""Request correlation and access logging middleware."""

from __future__ import annotations

import logging
import time
import uuid
from collections.abc import Awaitable, Callable

from fastapi import Request, Response
from starlette.datastructures import MutableHeaders
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import Message

logger = logging.getLogger(__name__)

RequestHandler = Callable[[Request], Awaitable[Response]]


class RequestContextMiddleware(BaseHTTPMiddleware):
    """Attach a bounded request ID and log one structured completion event."""

    def __init__(self, app: object, header_name: str) -> None:
        super().__init__(app)  # type: ignore[arg-type]
        self._header_name = header_name

    async def dispatch(self, request: Request, call_next: RequestHandler) -> Response:
        supplied = request.headers.get(self._header_name, "")
        request_id = supplied if 0 < len(supplied) <= 128 else str(uuid.uuid4())
        request.state.request_id = request_id
        started = time.monotonic()
        response = await call_next(request)
        response.headers[self._header_name] = request_id
        logger.info(
            "request completed",
            extra={
                "request_id": request_id,
                "method": request.method,
                "path": request.url.path,
                "status": response.status_code,
                "duration_ms": round((time.monotonic() - started) * 1000, 3),
            },
        )
        return response


class HeadAsGetMiddleware(BaseHTTPMiddleware):
    """Serve HEAD for every GET route (deep + stub) with an empty body.

    FastAPI ``add_api_route(methods=['GET'])`` and some router setups omit HEAD;
    contract matrix probes expect synthetic HEAD on each GET path.
    """

    async def dispatch(self, request: Request, call_next: RequestHandler) -> Response:
        if request.method != "HEAD":
            return await call_next(request)

        # Replay as GET, then strip the body while preserving status/headers.
        request.scope["method"] = "GET"
        response = await call_next(request)

        body = bytearray()
        async for chunk in response.body_iterator:
            if isinstance(chunk, str):
                body.extend(chunk.encode(response.charset or "utf-8"))
            else:
                body.extend(chunk)

        headers = MutableHeaders(scope={"type": "http", "headers": []})
        for key, value in response.headers.items():
            if key.lower() in {"content-length", "content-type", "transfer-encoding"}:
                continue
            headers.append(key, value)
        headers["content-length"] = str(len(body))
        if response.media_type:
            headers["content-type"] = response.media_type

        async def _empty_receive() -> Message:
            return {"type": "http.request", "body": b"", "more_body": False}

        del _empty_receive
        return Response(
            content=b"",
            status_code=response.status_code,
            headers=headers,
            media_type=response.media_type,
        )