"""跨进程限流后端: CHS 六道闸 Lua 逐字移植(D3 双后端;M2 设计 §2)。

语义蓝本 `reference/CHSAnalyzer/app/coordination/{scripts,limiter}.py`:
读-判-占在单条 Lua 内原子完成;并发 = ZSET 租约(score=过期时刻 ms,惰性
清理 + PEXPIRE 双保险);RPM/TPM = 分钟固定窗口计数,窗口 id 用 **Redis
服务器时钟**(多进程口径统一);TPM 预扣 settle 多退少补且退款落 acquire
时的窗口(INCRBY 可为负,存储保留负值保证跨窗口结算正确)。

对 CHS 的已声明偏离(设计 §9 勘误): ① 限额 0 = 该闸不启用(M1 契约冻结;
CHS Lua 无此守卫,0 在其语义下是全拒),Lua 判据显式加 `limit > 0`;
② key 前缀 `cclimit:` → `pgw:limit:`;③ `LimiterError` → `GovernanceBackendError`;
④ `source_stats` 读侧 clamp ≥0(展示口径与内存版统一,存储不 clamp)。
契约量纲为秒,毫秒换算是本后端私事(ports.py 模块 docstring)。
"""

from __future__ import annotations

import asyncio
import uuid
from typing import TYPE_CHECKING

from loguru import logger
from redis.exceptions import RedisError

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

if TYPE_CHECKING:
    from redis.asyncio import Redis

_WINDOW_TTL_S = 120  # 分钟窗口兜底过期(CHS limiter.py:16)
_PROGRESS_TTL_S = 3600  # 远大于任何 stall_window,防进度键过期造成假停滞(CHS :21)

# KEYS: g_lease s_lease g_rpm s_rpm g_tpm s_tpm(rpm/tpm key 已含 window 后缀)
# ARGV: g_conc s_conc g_rpm s_rpm g_tpm s_tpm est lease_ttl_ms window_ttl_s lease_id
# 拒绝零副作用(仅清过期 lease);闸判据 CHS scripts.py:18-23 + 库版 limit>0 守卫
ACQUIRE = """
local t = redis.call('TIME')
local now = tonumber(t[1]) * 1000 + math.floor(tonumber(t[2]) / 1000)
redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, now)
redis.call('ZREMRANGEBYSCORE', KEYS[2], 0, now)
local gi = redis.call('ZCARD', KEYS[1])
local si = redis.call('ZCARD', KEYS[2])
local gr = tonumber(redis.call('GET', KEYS[3]) or '0')
local sr = tonumber(redis.call('GET', KEYS[4]) or '0')
local gt = tonumber(redis.call('GET', KEYS[5]) or '0')
local st = tonumber(redis.call('GET', KEYS[6]) or '0')
local est = tonumber(ARGV[7])
if tonumber(ARGV[1]) > 0 and gi >= tonumber(ARGV[1]) then return 0 end
if tonumber(ARGV[2]) > 0 and si >= tonumber(ARGV[2]) then return 0 end
if tonumber(ARGV[3]) > 0 and gr >= tonumber(ARGV[3]) then return 0 end
if tonumber(ARGV[4]) > 0 and sr >= tonumber(ARGV[4]) then return 0 end
if tonumber(ARGV[5]) > 0 and gt + est > tonumber(ARGV[5]) then return 0 end
if tonumber(ARGV[6]) > 0 and st + est > tonumber(ARGV[6]) then return 0 end
local expire_at = now + tonumber(ARGV[8])
redis.call('ZADD', KEYS[1], expire_at, ARGV[10])
redis.call('ZADD', KEYS[2], expire_at, ARGV[10])
redis.call('PEXPIRE', KEYS[1], ARGV[8])
redis.call('PEXPIRE', KEYS[2], ARGV[8])
redis.call('INCR', KEYS[3]); redis.call('EXPIRE', KEYS[3], ARGV[9])
redis.call('INCR', KEYS[4]); redis.call('EXPIRE', KEYS[4], ARGV[9])
redis.call('INCRBY', KEYS[5], est); redis.call('EXPIRE', KEYS[5], ARGV[9])
redis.call('INCRBY', KEYS[6], est); redis.call('EXPIRE', KEYS[6], ARGV[9])
return 1
"""

# KEYS: g_lease s_lease ; ARGV: lease_id(ZREM 幂等,CHS scripts.py:37-41)
RELEASE = """
redis.call('ZREM', KEYS[1], ARGV[1])
redis.call('ZREM', KEYS[2], ARGV[1])
return 1
"""

# KEYS: g_tpm s_tpm ; ARGV: delta window_ttl_s(delta 可负=退款;落 acquire 窗口)
SETTLE = """
redis.call('INCRBY', KEYS[1], ARGV[1]); redis.call('EXPIRE', KEYS[1], ARGV[2])
redis.call('INCRBY', KEYS[2], ARGV[1]); redis.call('EXPIRE', KEYS[2], ARGV[2])
return 1
"""

# KEYS: s_lease s_rpm s_tpm(CHS scripts.py:51-58)
STATS = """
local t = redis.call('TIME')
local now = tonumber(t[1]) * 1000 + math.floor(tonumber(t[2]) / 1000)
redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, now)
return {redis.call('ZCARD', KEYS[1]),
        tonumber(redis.call('GET', KEYS[2]) or '0'),
        tonumber(redis.call('GET', KEYS[3]) or '0')}
"""

# KEYS: progress_key ; ARGV: ttl_s(CHS scripts.py:61-67)
PROGRESS_MARK = """
local t = redis.call('TIME')
local now = tonumber(t[1]) * 1000 + math.floor(tonumber(t[2]) / 1000)
redis.call('SET', KEYS[1], now)
redis.call('EXPIRE', KEYS[1], ARGV[1])
return 1
"""

# KEYS: progress_key(-1=从未成功;时钟回拨 clamp 0;CHS scripts.py:70-78)
PROGRESS_AGE = """
local v = redis.call('GET', KEYS[1])
if not v then return -1 end
local t = redis.call('TIME')
local now = tonumber(t[1]) * 1000 + math.floor(tonumber(t[2]) / 1000)
local d = now - tonumber(v)
if d < 0 then return 0 end
return d
"""


class _RedisPermit:
    """入场许可: 持有并发租约,负责释放与按 acquire 窗口结算 TPM(幂等)。"""

    def __init__(
        self, limiter: RedisLimiter, 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:
        """归还并发租约;释放侧 Redis 失败降级 warning(设计 §10),不冒泡。"""
        if self._released:
            return
        self._released = True
        try:
            await self._limiter._release_lease(self._source, self._lease_id)
        except GovernanceBackendError as exc:
            logger.warning("permit release 降级(租约将由 TTL 回收): {}", exc)

    async def settle(self, actual_tokens: int) -> None:
        """按实际 usage 结算 TPM 差额;释放侧失败降级 warning,不冒泡。

        有意偏离 CHS(limiter.py:62 成功后才置位): 本库释放侧失败被降级
        吞掉、调用方不重试,故尝试前置位——与内存版 flag 语义对齐,幂等
        性不受影响(2026-07-21 verifier M4 声明)。
        """
        if self._settled:
            return
        self._settled = True
        delta = actual_tokens - self._est
        if delta == 0:
            return
        try:
            await self._limiter._settle_tpm(self._source, delta, self._window)
        except GovernanceBackendError as exc:
            logger.warning("permit settle 降级(窗口计数将随 TTL 过期): {}", exc)


class RedisLimiter:
    """六道闸的跨进程实现;同 scope 的多进程 worker 天然共享全局限额。"""

    def __init__(
        self,
        *,
        scope: str,
        sources: dict[str, SourceConfig],
        global_limits: GlobalLimits,
        redis: Redis,
        lease_ttl_s: float = 1500.0,
        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.lower()
        self._sources = dict(sources)
        self._global = global_limits
        self._lease_ttl_ms = int(lease_ttl_s * 1000)
        self._poll_interval_s = poll_interval_s
        self._redis = redis
        self._owns_client = False
        self._acquire_lua = redis.register_script(ACQUIRE)
        self._release_lua = redis.register_script(RELEASE)
        self._settle_lua = redis.register_script(SETTLE)
        self._stats_lua = redis.register_script(STATS)
        self._progress_mark_lua = redis.register_script(PROGRESS_MARK)
        self._progress_age_lua = redis.register_script(PROGRESS_AGE)

    @classmethod
    def from_url(cls, url: str, **kwargs) -> RedisLimiter:
        """自建并持有 Redis 客户端(aclose 时代关);共享后端请直接注入 redis。"""
        import redis.asyncio as aioredis

        limiter = cls(redis=aioredis.from_url(url), **kwargs)
        limiter._owns_client = True
        return limiter

    # —— key 布局(CHS limiter.py:100-115,前缀改 pgw:limit:)——

    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 _lease_keys(self, source_key: str) -> tuple[str, str]:
        return (
            f"pgw:limit:GLOBAL:{self._scope}:lease",
            f"pgw:limit:{self._scope}:{source_key}:lease",
        )

    def _window_keys(self, source_key: str, window: int) -> dict[str, str]:
        gl = f"GLOBAL:{self._scope}"
        return {
            "g_rpm": f"pgw:limit:{gl}:rpm:{window}",
            "s_rpm": f"pgw:limit:{self._scope}:{source_key}:rpm:{window}",
            "g_tpm": f"pgw:limit:{gl}:tpm:{window}",
            "s_tpm": f"pgw:limit:{self._scope}:{source_key}:tpm:{window}",
        }

    def _progress_key(self) -> str:
        return f"pgw:limit:GLOBAL:{self._scope}:progress_ms"

    async def _window_id(self) -> int:
        """Redis 服务器时钟生成分钟窗口 id(CHS limiter.py:95-98)。"""
        sec, _ = await self._redis.time()
        return int(sec) // 60

    # —— 契约方法 ——

    async def try_acquire(self, source_key: str, est_tokens: int) -> _RedisPermit | None:
        """非阻塞准入: 六闸单 Lua 原子判定;拒绝零副作用;Redis 失败报错不放行。"""
        cfg = self._cfg(source_key)
        lease_id = uuid.uuid4().hex
        try:
            window = await self._window_id()
            gl, sl = self._lease_keys(source_key)
            wk = self._window_keys(source_key, window)
            ok = await self._acquire_lua(
                keys=[gl, sl, wk["g_rpm"], wk["s_rpm"], wk["g_tpm"], wk["s_tpm"]],
                args=[
                    self._global.max_concurrency,
                    cfg.max_concurrency,
                    self._global.rpm,
                    cfg.rpm,
                    self._global.tpm,
                    cfg.tpm,
                    est_tokens,
                    self._lease_ttl_ms,
                    _WINDOW_TTL_S,
                    lease_id,
                ],
            )
        except RedisError as exc:
            raise GovernanceBackendError(f"限流后端 try_acquire 失败: {exc}", scope=self._scope) from exc
        if ok != 1:
            return None
        return _RedisPermit(self, source_key, lease_id, est_tokens, window)

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

    async def _release_lease(self, source_key: str, lease_id: str) -> None:
        gl, sl = self._lease_keys(source_key)
        try:
            await self._release_lua(keys=[gl, sl], args=[lease_id])
        except RedisError as exc:
            raise GovernanceBackendError(f"限流后端 release 失败: {exc}", scope=self._scope) from exc

    async def _settle_tpm(self, source_key: str, delta: int, window: int) -> None:
        wk = self._window_keys(source_key, window)
        try:
            await self._settle_lua(keys=[wk["g_tpm"], wk["s_tpm"]], args=[delta, _WINDOW_TTL_S])
        except RedisError as exc:
            raise GovernanceBackendError(f"限流后端 settle 失败: {exc}", scope=self._scope) from exc

    async def source_stats(self, source_key: str) -> SourceStats:
        """当前窗口快照;读侧 clamp ≥0(展示口径,存储保留负值)。"""
        self._cfg(source_key)
        try:
            window = await self._window_id()
            _, sl = self._lease_keys(source_key)
            wk = self._window_keys(source_key, window)
            res = await self._stats_lua(keys=[sl, wk["s_rpm"], wk["s_tpm"]])
        except RedisError as exc:
            raise GovernanceBackendError(f"限流后端 source_stats 失败: {exc}", scope=self._scope) from exc
        return SourceStats(
            inflight=int(res[0]),
            rpm_used=max(0, int(res[1])),
            tpm_used=max(0, int(res[2])),
        )

    async def mark_progress(self) -> None:
        """记录 scope 级"最近出餐"时刻;失败报错,降级责任在中间件记账侧(T5)。"""
        try:
            await self._progress_mark_lua(keys=[self._progress_key()], args=[_PROGRESS_TTL_S])
        except RedisError as exc:
            raise GovernanceBackendError(f"限流后端 mark_progress 失败: {exc}", scope=self._scope) from exc

    async def progress_age_s(self) -> float:
        """距上次全局成功的秒数;仅键缺失(-1)= 从未进展 → inf(CHS limiter.py:208)。"""
        try:
            res = await self._progress_age_lua(keys=[self._progress_key()])
        except RedisError as exc:
            raise GovernanceBackendError(f"限流后端 progress_age_s 失败: {exc}", scope=self._scope) from exc
        return float("inf") if int(res) == -1 else int(res) / 1000.0

    async def aclose(self) -> None:
        """幂等释放自建客户端;注入的客户端归注入方管理。"""
        if self._owns_client:
            self._owns_client = False
            await self._redis.aclose()