"""API Key 轮换环:多 key 自动切换,遇到限流(429)自动跳过已冷却 key。

设计:
  - 进程内纯内存状态(限速是短时状态,无需持久化;D7: 多 worker 部署时各进程
    独立计数、不共享,上游限速额度应按 worker 数分摊,或前置网关统一管理)
  - 单 key 场景零开销:直接透传
  - 多 key 场景:429 时把当前 key 标记冷却(默认 60s),轮转到下一个可用 key
  - 全部 key 都在冷却时:使用最早冷却的那个 key 并等待(退化到单 key 重试)
"""
from __future__ import annotations

import logging
import time

logger = logging.getLogger(__name__)

DEFAULT_COOLDOWN = 60.0  # 单个 key 被限流后的冷却时间(秒)


class KeyRing:
    """多 key 轮换环。在 async 单线程事件循环下无需加锁。"""

    def __init__(self, keys: list[str], cooldown_seconds: float = DEFAULT_COOLDOWN) -> None:
        self._keys: list[str] = [k.strip() for k in keys if k and k.strip()]
        self._cooldown = cooldown_seconds
        # key → 冷却过期时间戳(时刻);不在表中表示可用
        self._blocked_until: dict[str, float] = {}
        self._index = 0  # 当前轮转位置

    @property
    def has_multiple(self) -> bool:
        return len(self._keys) > 1

    @property
    def all_keys(self) -> list[str]:
        return list(self._keys)

    @property
    def active_key(self) -> str | None:
        """当前指向的 key(即使正在冷却也返回,用于上报)。"""
        if not self._keys:
            return None
        return self._keys[self._index]

    def get(self) -> str | None:
        """获取一个可用 key:优先选不在冷却中的;全部冷却则选最早过期的。"""
        if not self._keys:
            return None
        if len(self._keys) == 1:
            return self._keys[0]

        now = time.monotonic()
        n = len(self._keys)
        # 从当前 index 开始找一圈,找一个可用的
        for offset in range(n):
            idx = (self._index + offset) % n
            key = self._keys[idx]
            expire = self._blocked_until.get(key, 0.0)
            if now >= expire:
                # 可用:把指针移到这里
                self._index = idx
                # 清理已过期的冷却记录
                if key in self._blocked_until:
                    del self._blocked_until[key]
                return key

        # 全部在冷却中:选最早过期的那个,并等待到它过期
        earliest_key = min(self._keys, key=lambda k: self._blocked_until.get(k, 0.0))
        self._index = self._keys.index(earliest_key)
        return earliest_key

    def report_rate_limited(self, key: str | None = None) -> str | None:
        """上报某个 key 被限流(429)。默认是当前 key。返回切换后的新 key。"""
        target = key or self.active_key
        if target and len(self._keys) > 1:
            until = time.monotonic() + self._cooldown
            self._blocked_until[target] = until
            logger.warning(
                "bzzoiro key 被限流(429),冷却 %.0fs: %s", self._cooldown, _mask(target),
            )
        # 轮转到下一个(即使只有一个 key 也做一次 get,保持行为一致)
        return self.get()

    def wait_if_all_blocked(self) -> float:
        """如果所有 key 都在冷却中,返回需要等待的秒数;否则返回 0。"""
        if len(self._keys) <= 1:
            return 0.0
        now = time.monotonic()
        remaining = [self._blocked_until.get(k, 0.0) - now for k in self._keys]
        if all(r > -0.001 for r in remaining) and any(r > 0.001 for r in remaining):
            # 全部仍在冷却中
            return max(remaining)
        return 0.0

    def stats(self) -> dict:
        """当前 key 环状态(用于管理后台展示)。"""
        now = time.monotonic()
        return {
            "total": len(self._keys),
            "keys": [
                {
                    "masked": _mask(k),
                    "blocked_remaining": max(0.0, round(self._blocked_until.get(k, 0.0) - now, 1)),
                }
                for k in self._keys
            ],
            "active_index": self._index,
        }


def _mask(key: str) -> str:
    """脱敏:只显示前 4 位和后 4 位。"""
    if len(key) <= 10:
        return key[:2] + "***"
    return key[:4] + "..." + key[-4:]


# ── 全局单例(进程级,按 base URL 隔离) ──────────────────────────
_RINGS: dict[str, KeyRing] = {}


def parse_keys(value: str | None) -> list[str]:
    """解析 key 配置值:支持逗号、分号、换行分隔的多个 key。"""
    if not value:
        return []
    # 统一替换分隔符为逗号后拆分
    normalized = value.replace("\n", ",").replace(";", ",")
    return [k.strip() for k in normalized.split(",") if k.strip()]


def get_key_ring(base: str, raw_keys: str | None, cooldown_seconds: float = DEFAULT_COOLDOWN) -> KeyRing:
    """获取(或创建)某 base URL 对应的 KeyRing。"""
    key = base
    ring = _RINGS.get(key)
    parsed = parse_keys(raw_keys)
    if ring is None:
        ring = KeyRing(parsed, cooldown_seconds)
        _RINGS[key] = ring
    else:
        # 热更新 key 列表(增删 key 无需重启)
        if set(ring.all_keys) != set(parsed):
            ring._keys = parsed
            ring._blocked_until.clear()
            ring._index = 0
        ring._cooldown = cooldown_seconds
    return ring