"""Correlation ID middleware -- async-safe context propagation for distributed tracing.

Extracts or generates a correlation ID from ``X-Correlation-ID`` header and
stores it in async-safe context storage (contextvars). The ID is returned in
the response header and can be used for log correlation across services.

This module is part of T-12 (World-Class Logging & Monitoring) compliance.
"""

from __future__ import annotations

import contextvars
import logging
import uuid
from typing import TYPE_CHECKING

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint

if TYPE_CHECKING:
    from starlette.requests import Request
    from starlette.responses import Response

logger = logging.getLogger(__name__)

# Async-safe storage for correlation ID
_correlation_id_context: contextvars.ContextVar[str] = contextvars.ContextVar("correlation_id", default="")


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def get_correlation_id() -> str:
    """Get the current correlation ID from context.

    Returns
    -------
    str
        The correlation ID for the current request, or empty string if not set.
    """
    return _correlation_id_context.get()


def set_correlation_id(correlation_id: str) -> None:
    """Set the correlation ID in async context.

    Parameters
    ----------
    correlation_id:
        The correlation ID to set.
    """
    _correlation_id_context.set(correlation_id)


# ---------------------------------------------------------------------------
# Middleware
# ---------------------------------------------------------------------------

CORRELATION_ID_HEADER = "X-Correlation-ID"


class CorrelationIdMiddleware(BaseHTTPMiddleware):
    """Extract or generate a correlation ID and store in async-safe context.

    The correlation ID is returned in the response header and can be accessed
    via :func:`get_correlation_id` from any async function in the request context.
    """

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        # Extract or generate correlation ID
        correlation_id = request.headers.get(CORRELATION_ID_HEADER)
        if not correlation_id:
            correlation_id = uuid.uuid4().hex

        # Store in async-safe context
        set_correlation_id(correlation_id)

        # Also store on request state for compatibility
        request.state.correlation_id = correlation_id

        # Process request
        response = await call_next(request)

        # Return correlation ID in response header
        response.headers[CORRELATION_ID_HEADER] = correlation_id

        return response
