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

# ContextVar to hold request_id globally in the thread/coroutine context
request_id_var = contextvars.ContextVar("request_id", default="")

def get_request_id() -> str:
    return request_id_var.get()

class RequestTraceMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # Generate new trace ID
        req_id = str(ulid.ULID())
        
        # Set context variable
        token = request_id_var.set(req_id)
        
        # Store in state for easy route access
        request.state.request_id = req_id
        
        try:
            response = await call_next(request)
            # Append to response headers
            response.headers["X-Request-ID"] = req_id
            return response
        finally:
            # Reset context variable to prevent leakage
            request_id_var.reset(token)