"""
请求追踪中间件

为每个请求生成唯一 ID，记录请求日志，添加响应头
""" 
import uuid
import time
import logging
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
from src.utils.logging import set_log_context, clear_log_context

logger = logging.getLogger(__name__)


class RequestTrackingMiddleware(BaseHTTPMiddleware):
    """请求追踪中间件
    
    功能：
    1. 为每个请求生成唯一 request_id (correlation ID)
    2. 记录请求开始和结束时间
    3. 在响应头中添加 request_id 和响应时间
    4. 记录请求日志（包含耗时、状态码等）
    5. 设置日志上下文，使所有日志自动包含 request_id
    """ 
    async def dispatch(self, request: Request, call_next):
        # 生成或获取请求 ID (correlation ID)
        request_id = request.headers.get("X-Request-ID") or str(uuid.uuid4())
        
        # 将 request_id 存储到 request.state 中，供后续使用
        request.state.request_id = request_id
        
        # 设置日志上下文，使所有后续日志自动包含 request_id
        set_log_context(request_id=request_id)
        
        # 记录请求开始时间
        start_time = time.time()
        
        # 记录请求开始日志
        logger.info(
            f"Request started: {request.method} {request.url.path}",
            extra={
                'extra_fields': {
                    "method": request.method,
                    "path": request.url.path,
                    "query_params": str(request.query_params),
                    "client_host": request.client.host if request.client else None,
                    "user_agent": request.headers.get("user-agent")
                }
            }
        )
        
        # 处理请求
        try:
            response: Response = await call_next(request)
        except Exception as e:
            # 记录异常
            duration = time.time() - start_time
            logger.error(
                f"Request failed: {request.method} {request.url.path}",
                extra={
                    'extra_fields': {
                        "method": request.method,
                        "path": request.url.path,
                        "duration": duration,
                        "error": str(e)
                    }
                },
                exc_info=True
            )
            # 清除日志上下文
            clear_log_context()
            raise
        
        # 计算请求耗时
        duration = time.time() - start_time
        
        # 添加响应头
        response.headers["X-Request-ID"] = request_id
        response.headers["X-Response-Time"] = f"{duration:.3f}s"
        
        # 记录请求完成日志
        log_level = logging.INFO
        if response.status_code >= 500:
            log_level = logging.ERROR
        elif response.status_code >= 400:
            log_level = logging.WARNING
        
        logger.log(
            log_level,
            f"Request completed: {request.method} {request.url.path} - {response.status_code}",
            extra={
                'extra_fields': {
                    "method": request.method,
                    "path": request.url.path,
                    "status_code": response.status_code,
                    "duration": duration,
                    "response_size": response.headers.get("content-length")
                }
            }
        )
        
        # 清除日志上下文
        clear_log_context()
        
        return response


def setup_request_tracking(app):
    """注册请求追踪中间件"""
    app.add_middleware(RequestTrackingMiddleware)
