import re
import uuid

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

from ..logging_config import request_id_var

_UUID_RE = re.compile(r"^[0-9a-fA-F-]{8,64}$")


class RequestIdMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        incoming = request.headers.get("X-Request-Id")
        rid = incoming if incoming and _UUID_RE.match(incoming) else str(uuid.uuid4())
        request.state.request_id = rid
        token = request_id_var.set(rid)
        try:
            response = await call_next(request)
        finally:
            request_id_var.reset(token)
        response.headers["X-Request-Id"] = rid
        return response