import os
from typing import List, Optional

from prometheus_client import Counter, Histogram
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request


class RequestTimeCostMiddleware(BaseHTTPMiddleware):
    def __init__(
        self,
        app,
        buckets: Optional[List[float]] = None,
        ignore_paths: Optional[List[str]] = None,
    ):
        """
        :param buckets: 默认[0.1, 0.5, 1, 3, 5]
        :param ignore_paths: 默认不屏蔽，但是建议用户屏蔽/metrics
        """
        super().__init__(app)

        default_buckets = [0.1, 0.5, 1, 3, 5]
        self.buckets = buckets if buckets is not None else default_buckets
        self.ignore_paths = ignore_paths if ignore_paths is not None else []

        self.request_duration = Histogram(
            "http_request_duration_seconds",
            "Request duration in seconds",
            ["uri", "pid"],
            buckets=self.buckets,
        )

        self.error_requests = Counter(
            "http_error_requests", "error HTTP requests", ["uri", "status"]
        )

    async def dispatch(self, request: Request, call_next):
        # 规范化路径：去掉末尾的斜杠（根路径除外）
        path = request.url.path
        normalized_path = path.rstrip("/") if path != "/" else path

        if normalized_path in self.ignore_paths:
            return await call_next(request)
        else:
            with self.request_duration.labels(
                uri=normalized_path, pid=os.getpid()
            ).time():
                response = await call_next(request)

                # 记录错误请求
                if response.status_code >= 400:
                    self.error_requests.labels(
                        uri=normalized_path, status=response.status_code
                    ).inc()

                return response
