"""Correlation ID middleware for distributed tracing.

This middleware generates or extracts a correlation ID from each request,
stores it in request state, and includes it in error responses.
This enables correlating logs across frontend and backend for a single
user action or request flow.

Correlation IDs flow through the request lifecycle:
1. Frontend generates/passes via `X-Correlation-ID` header
2. Middleware extracts or generates a UUID4
3. Stores on request.state for use by error handlers and log filters
4. Error responses include the correlation ID for client-side correlation

Processing order
-----------------
This middleware must be the outermost in the security-critical chain so it
executes first on incoming requests (outermost = first to see request,
last to see response). In the required chain:

    CorrelationIdMiddleware → CsrfMiddleware → RateLimitMiddleware

The registration order in ``main.py`` must be:
    RateLimitMiddleware, CsrfMiddleware, CorrelationIdMiddleware
(last registered = outermost in Starlette's reverse application).
"""

from __future__ import annotations

from app.utils.logging_compat import get_logger
import uuid
from typing import TYPE_CHECKING

from starlette.middleware.base import BaseHTTPMiddleware

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from starlette.requests import Request
    from starlette.responses import Response as StarletteResponse

log = get_logger(__name__)

# Standard header name for correlation IDs (follows W3C Trace Context conventions)
_CORRELATION_ID_HEADER: str = "X-Correlation-ID"

# Key name for storing correlation ID in request state
CORRELATION_ID_CONTEXT_KEY: str = "correlation_id"


class CorrelationIdMiddleware(BaseHTTPMiddleware):
    """Extract or generate correlation ID and store on request state.

    For each request, this middleware:
    1. Checks for `X-Correlation-ID` header (trusted from frontend)
    2. Generates a new UUID4 if header not present
    3. Stores on request.state for use by error handlers and log filters

    The correlation ID enables tracing a single user action or request flow
    across both frontend and backend systems using structured logs.
    """

    async def dispatch(
        self,
        request: Request,
        call_next: Callable[[Request], Awaitable[StarletteResponse]],
    ) -> StarletteResponse:
        """Intercept requests to extract or generate correlation ID.

        Args:
            request: The incoming HTTP request.
            call_next: The next middleware / router handler.

        Returns:
            The response from the next middleware / router, with correlation ID
            in the request state for use by exception handlers.
        """
        # Extract correlation ID from request header, or generate a new one
        correlation_id: str = request.headers.get(
            _CORRELATION_ID_HEADER,
            str(uuid.uuid4()),
        )

        # Store on request.state for use by exception handlers
        request.state.correlation_id = correlation_id

        log.debug(
            "request_received",
            extra={"method": request.method, "path": request.url.path},
        )

        response: StarletteResponse = await call_next(request)

        # Add correlation ID to response header so frontend can correlate errors
        response.headers[_CORRELATION_ID_HEADER] = correlation_id

        return response
