from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from app.core.middleware.audit_context import request_id_var, user_id_var, ip_address_var, user_agent_var
import ulid

class AuditMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # 1. Set IP Address
        x_forwarded_for = request.headers.get("x-forwarded-for")
        ip = x_forwarded_for.split(",")[0].strip() if x_forwarded_for else (request.client.host if request.client else "127.0.0.1")
        
        token_ip = ip_address_var.set(ip)
        token_ua = user_agent_var.set(request.headers.get("user-agent", ""))
        
        # 2. Set request ID
        req_id = getattr(request.state, "request_id", None)
        if not req_id:
            req_id = str(ulid.ULID())
            request.state.request_id = req_id
        token_req = request_id_var.set(req_id)
        
        # Default user_id to None
        token_user = user_id_var.set(None)
        
        try:
            response = await call_next(request)
            return response
        finally:
            ip_address_var.reset(token_ip)
            user_agent_var.reset(token_ua)
            request_id_var.reset(token_req)
            user_id_var.reset(token_user)