"""进程内限流后端: 与 Redis 版同一契约的六道闸实现(D3 双后端)。

语义蓝本 CHS `app/coordination/limiter.py`: 并发 = 带 TTL 的租约(持有者
死亡后过期回收);RPM/TPM = 分钟**固定窗口**计数(`int(now/60)`,与 CHS
Lua 同款口径);TPM 入场按 est 预扣,settle
按实际结算多退少补且退款落 acquire 时的窗口。检查-占用在单次同步段内完成
(无 await 穿插),单事件循环下天然原子;本实现不跨进程,是单进程部署的
正确答案(多 worker 用 M2 Redis 后端)。
"""

from __future__ import annotations

import asyncio
import time
import uuid
from typing import TYPE_CHECKING

from polygateway.errors import SourceNotConfiguredError
from polygateway.types import GlobalLimits, SourceConfig, SourceStats

if TYPE_CHECKING:
    from collections.abc import Callable

_GLOBAL = "__global__"
_WINDOW_S = 60.0


class _MemoryPermit:
    """入场许可;release/settle 幂等(CHS _RedisPermit 同款 flag 语义)。"""

    def __init__(
        self, limiter: InMemoryLimiter, source: str, lease_id: str, est: int, window: int
    ) -> None:
        self._limiter = limiter
        self._source = source
        self._lease_id = lease_id
        self._est = est
        self._window = window
        self._released = False
        self._settled = False

    async def release(self) -> None:
        if self._released:
            return
        self._released = True
        self._limiter._release_lease(self._source, self._lease_id)

    async def settle(self, actual_tokens: int) -> None:
        if self._settled:
            return
        self._settled = True
        delta = actual_tokens - self._est
        if delta:
            self._limiter._settle_tpm(self._source, delta, self._window)


class InMemoryLimiter:
    """六道闸: 全局/单源 × 并发/RPM/TPM;限额 0 = 该闸不启用。"""

    def __init__(
        self,
        *,
        scope: str,
        sources: dict[str, SourceConfig],
        global_limits: GlobalLimits,
        lease_ttl_s: float = 1500.0,
        now: Callable[[], float] = time.monotonic,
        sleep: Callable[[float], object] = asyncio.sleep,
        poll_interval_s: float = 0.05,
    ) -> None:
        if lease_ttl_s <= 0 or poll_interval_s <= 0:
            raise ValueError("lease_ttl_s 与 poll_interval_s 必须 > 0")
        self._scope = scope
        self._sources = dict(sources)
        self._global = global_limits
        self._lease_ttl_s = lease_ttl_s
        self._now = now
        self._sleep = sleep
        self._poll_interval_s = poll_interval_s
        # source → lease_id → 过期时刻
        self._leases: dict[str, dict[str, float]] = {name: {} for name in self._sources}
        # (source|__global__, window_id) → 计数
        self._rpm: dict[tuple[str, int], int] = {}
        self._tpm: dict[tuple[str, int], int] = {}
        self._progress_at: float | None = None

    # —— 内部原语(同步,单事件循环下原子)——

    def _cfg(self, source_key: str) -> SourceConfig:
        cfg = self._sources.get(source_key)
        if cfg is None:
            raise SourceNotConfiguredError(f"未知源 {source_key!r}(scope={self._scope})")
        return cfg

    def _window(self) -> int:
        return int(self._now() / _WINDOW_S)

    def _purge_leases(self) -> None:
        now = self._now()
        for leases in self._leases.values():
            expired = [lid for lid, expiry in leases.items() if expiry <= now]
            for lid in expired:
                del leases[lid]

    def _inflight(self, source_key: str) -> int:
        return len(self._leases[source_key])

    def _inflight_total(self) -> int:
        return sum(len(leases) for leases in self._leases.values())

    def _release_lease(self, source_key: str, lease_id: str) -> None:
        self._leases[source_key].pop(lease_id, None)

    def _settle_tpm(self, source_key: str, delta: int, window: int) -> None:
        # 退款/补记落 acquire 时的窗口,计数下限 0(CHS Lua 同款)
        for key in (source_key, _GLOBAL):
            slot = (key, window)
            self._tpm[slot] = max(0, self._tpm.get(slot, 0) + delta)

    def _gates_pass(self, cfg: SourceConfig, est_tokens: int, window: int) -> bool:
        checks = (
            (self._global.max_concurrency, self._inflight_total() + 1),
            (cfg.max_concurrency, self._inflight(cfg.name) + 1),
            (self._global.rpm, self._rpm.get((_GLOBAL, window), 0) + 1),
            (cfg.rpm, self._rpm.get((cfg.name, window), 0) + 1),
            (self._global.tpm, self._tpm.get((_GLOBAL, window), 0) + est_tokens),
            (cfg.tpm, self._tpm.get((cfg.name, window), 0) + est_tokens),
        )
        return all(limit <= 0 or would_be <= limit for limit, would_be in checks)

    # —— 契约方法 ——

    async def try_acquire(self, source_key: str, est_tokens: int) -> _MemoryPermit | None:
        """非阻塞准入: 六闸全过才占用;任一满则返回 None 且零副作用。"""
        cfg = self._cfg(source_key)
        self._purge_leases()
        window = self._window()
        if not self._gates_pass(cfg, est_tokens, window):
            return None
        lease_id = uuid.uuid4().hex
        self._leases[source_key][lease_id] = self._now() + self._lease_ttl_s
        for key in (source_key, _GLOBAL):
            self._rpm[(key, window)] = self._rpm.get((key, window), 0) + 1
            if est_tokens:
                self._tpm[(key, window)] = self._tpm.get((key, window), 0) + est_tokens
        return _MemoryPermit(self, source_key, lease_id, est_tokens, window)

    async def acquire(self, source_key: str, est_tokens: int) -> _MemoryPermit:
        """阻塞准入: 轮询直到拿到 permit;sleep 可被取消(取消穿透铁律)。"""
        while True:
            permit = await self.try_acquire(source_key, est_tokens)
            if permit is not None:
                return permit
            await self._sleep(self._poll_interval_s)

    async def source_stats(self, source_key: str) -> SourceStats:
        self._cfg(source_key)
        self._purge_leases()
        window = self._window()
        return SourceStats(
            inflight=self._inflight(source_key),
            rpm_used=self._rpm.get((source_key, window), 0),
            tpm_used=self._tpm.get((source_key, window), 0),
        )

    async def mark_progress(self) -> None:
        """记录"最近一次出餐"时刻,供背压 stall 判定(M2)读取。"""
        self._progress_at = self._now()

    async def progress_age_s(self) -> float:
        if self._progress_at is None:
            return float("inf")
        return self._now() - self._progress_at