"""
中间件定义
"""
import time
import uuid
from typing import Callable

from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware

from app.core.logger import logger


class RequestIDMiddleware(BaseHTTPMiddleware):
    """请求ID中间件"""

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        # 生成请求ID
        request_id = str(uuid.uuid4())

        # 将请求ID添加到request状态
        request.state.request_id = request_id

        # 记录请求开始
        start_time = time.time()

        # 处理请求
        response = await call_next(request)

        # 计算处理时间
        process_time = time.time() - start_time

        # 添加响应头
        response.headers["X-Request-ID"] = request_id
        response.headers["X-Process-Time"] = str(process_time)

        # 记录请求日志
        logger.info(
            "HTTP请求",
            method=request.method,
            url=str(request.url),
            status_code=response.status_code,
            process_time=process_time,
            request_id=request_id,
        )

        return response


class GlobalContextMiddleware(BaseHTTPMiddleware):
    """全局上下文中间件"""

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        # 设置追踪ID（用于分布式追踪）
        trace_id = request.headers.get("X-Trace-ID", str(uuid.uuid4()))
        request.state.trace_id = trace_id

        # 处理请求
        response = await call_next(request)

        # 添加追踪ID到响应头
        response.headers["X-Trace-ID"] = trace_id

        return response