from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp
import time
from typing import Callable
from ..utils import get_logger

logger = get_logger(__name__)


class PerformanceMiddleware(BaseHTTPMiddleware):
    """性能监控中间件"""
    
    # 类变量，用于存储实例引用
    _instance = None

    def __init__(self, app: ASGIApp):
        super().__init__(app)
        self.logger = get_logger("performance")
        self.total_requests = 0
        self.requests_per_status = {}
        self.total_response_time = 0
        
        # 保存实例引用
        PerformanceMiddleware._instance = self

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        """处理请求并添加性能头"""
        start_time = time.time()

        try:
            response = await call_next(request)
            response_time = time.time() - start_time

            # 添加性能头
            response.headers["X-Response-Time"] = f"{response_time:.4f}"
            
            # 更新请求计数
            self.total_requests += 1
            status_code = response.status_code
            if status_code not in self.requests_per_status:
                self.requests_per_status[status_code] = 0
            self.requests_per_status[status_code] += 1
            
            # 更新总响应时间
            self.total_response_time += response_time

            return response

        except Exception as e:
            raise
    
    def get_stats(self):
        """获取性能统计信息"""
        return {
            "total_requests": self.total_requests,
            "requests_per_status": self.requests_per_status,
            "total_response_time": self.total_response_time,
            "average_response_time": self.total_response_time / self.total_requests if self.total_requests > 0 else 0
        }
    
    @classmethod
    def get_instance(cls):
        """获取中间件实例"""
        return cls._instance