"""请求审计上下文中间件。"""

from __future__ import annotations

from uuid import uuid4

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

from ..utils.audited_write import AuditRequestContext, normalize_caller, reset_audit_context, set_audit_context


class AuditContextMiddleware(BaseHTTPMiddleware):
    """为每个请求注入 trace/caller/path 审计上下文。"""

    @staticmethod
    def _infer_caller(request: Request) -> str:
        caller_header = request.headers.get("X-Caller")
        if caller_header:
            return normalize_caller(caller_header)
        path = str(request.url.path or "")
        if path.startswith(("/webhooks", "/feishu", "/wecom")):
            return "ai"
        if request.headers.get("Authorization"):
            return "user"
        return "system"

    async def dispatch(self, request: Request, call_next) -> Response:
        trace_id = str(request.headers.get("X-Trace-Id") or f"req-{uuid4().hex[:24]}")[:64]
        context = AuditRequestContext(
            trace_id=trace_id,
            interface_path=str(request.url.path or "")[:255] or "/",
            caller=self._infer_caller(request),
            method=str(request.method or "").upper()[:16],
        )
        token = set_audit_context(context)
        request.state.trace_id = trace_id
        try:
            response = await call_next(request)
        finally:
            reset_audit_context(token)
        response.headers["X-Trace-Id"] = trace_id
        return response