"""Async correlation ID context variable and helpers."""

import contextvars
import uuid

from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware

CORRELATION_ID: contextvars.ContextVar[str] = contextvars.ContextVar("correlation_id")


def get_correlation_id() -> str:
    """Return the current correlation ID or generate a new UUID."""
    try:
        return CORRELATION_ID.get()
    except LookupError:
        return str(uuid.uuid4())


class CorrelationIdMiddleware(BaseHTTPMiddleware):
    """Set correlation ID from X-Request-ID header or generate a new UUID."""

    async def dispatch(self, request: Request, call_next):
        request_id = request.headers.get("X-Request-ID")
        correlation_id = request_id or str(uuid.uuid4())
        token = CORRELATION_ID.set(correlation_id)
        try:
            response = await call_next(request)
            response.headers["X-Request-ID"] = correlation_id
            return response
        finally:
            CORRELATION_ID.reset(token)