"""在途请求计数器与中间件，用于优雅关停时排空在途请求。

``InflightRequestTracker`` 通过 ``increment`` / ``decrement`` 维护在途请求
计数，计数归零时 ``wait_drained`` 立即返回；超时未归零则抛 ``TimeoutError``，
由 ``HostShutdown._waitForInflightRequests`` 捕获并记录告警后强制进入下一步
（§11.9）。

``InflightTrackingMiddleware`` 在每个 HTTP 请求进入时 ``increment``、结束时
``decrement``（``finally`` 保证异常路径也执行），确保计数准确。
"""

from __future__ import annotations

import asyncio

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

__all__ = ["InflightRequestTracker", "InflightTrackingMiddleware"]


class InflightRequestTracker:
    """在途请求计数器，用于优雅关停时排空在途请求。

    关键不变量：
        - ``_count`` 归零时 ``_zero_event`` 被 set，``wait_drained`` 立即返回。
        - ``decrement`` 防御性归零，避免计数变为负数。
    """

    def __init__(self) -> None:
        self._count = 0
        self._zero_event = asyncio.Event()
        self._zero_event.set()  # 初始为 0，已触发

    @property
    def count(self) -> int:
        """当前在途请求数。"""
        return self._count

    def increment(self) -> None:
        """在途请求计数 +1，清除归零事件。"""
        self._count += 1
        self._zero_event.clear()

    def decrement(self) -> None:
        """在途请求计数 -1，归零时触发归零事件。"""
        self._count -= 1
        if self._count <= 0:
            self._count = 0
            self._zero_event.set()

    async def wait_drained(self, timeout: float) -> None:
        """等待在途请求排空。

        参数：
            timeout: 等待超时（秒），超时后抛 ``TimeoutError``。

        Raises:
            TimeoutError: 超时未排空。
        """
        await asyncio.wait_for(self._zero_event.wait(), timeout=timeout)


class InflightTrackingMiddleware(BaseHTTPMiddleware):
    """在途请求计数中间件。

    每个请求进入时 ``increment``、结束时 ``decrement``（``finally`` 保证
    异常路径也执行）。注册时应最后注册（LIFO 顺序中最先执行），确保覆盖
    所有后续中间件与路由处理。
    """

    async def dispatch(self, request: Request, call_next):
        tracker: InflightRequestTracker = request.app.state.inflight_tracker
        tracker.increment()
        try:
            return await call_next(request)
        finally:
            tracker.decrement()