"""
NVIDIA Sidecar 限流代理 — 令牌桶 + 网关识别模块 (§3.2)

从 BIZ-26 rate_limiter.py 提取核心限流逻辑，去除多线程调度器、缓存管理等。
保留：Priority, TokenBucket, is_nvidia_gateway, normalize_gateway_name。
"""

from __future__ import annotations

import time
import threading
from enum import IntEnum
from typing import Any


# ---------------------------------------------------------------------------
# 优先级枚举
# ---------------------------------------------------------------------------

class Priority(IntEnum):
    """请求优先级（数值越小优先级越高）。"""
    URGENT = 1
    HIGH = 2
    NORMAL = 3
    LOW = 4


# ---------------------------------------------------------------------------
# NVIDIA 网关别名集
# ---------------------------------------------------------------------------

NVIDIA_GATEWAY_ALIASES: set[str] = {
    "nvidia",
    "nvidia-gateway",
    "nvidiavx",
    "nvidiavx18088980513",
}


def is_nvidia_gateway(value: str | None) -> bool:
    """判断给定网关名/模型全路径是否属于 NVIDIA 网关。

    Args:
        value: 网关名（如 ``"nvidia"``）或模型全路径前缀
               （如 ``"nvidia/deepseek-ai/deepseek-v4-pro"``）。
               None 时直接返回 False。

    Returns:
        True 当 value 的 provider 部分匹配已知 NVIDIA 别名。
    """
    if value is None:
        return False

    # 提取 provider 前缀：取 "/" 前第一个部分
    provider = value.split("/", 1)[0].lower().strip()
    return provider in NVIDIA_GATEWAY_ALIASES


def normalize_gateway_name(value: str | None) -> str | None:
    """规范化网关名：提取 provider 前缀并转为小写。

    Args:
        value: 网关名或模型全路径。None 时返回 None。

    Returns:
        provider 前缀的小写形式，或 None。
    """
    if value is None:
        return None
    return value.split("/", 1)[0].lower().strip()


# ---------------------------------------------------------------------------
# 令牌桶（线程安全）
# ---------------------------------------------------------------------------

class TokenBucket:
    """线程安全的令牌桶实现。

    支持固定速率令牌补充和消费，带有溢出保护和可选的阻塞等待。
    """

    def __init__(self, rate: float = 40 / 60, capacity: int = 40) -> None:
        """初始化令牌桶。

        Args:
            rate: 令牌补充速率（令牌/秒）。默认 40/60 ≈ 0.667 token/s（40 RPM）。
            capacity: 桶最大容量（令牌数）。默认 40。
        """
        self._rate: float = float(rate)
        self._capacity: int = int(capacity)
        self._tokens: float = float(capacity)  # 启动时桶满
        self._last_refill: float = time.monotonic()
        self._lock: threading.Lock = threading.Lock()

    # ---- 内部方法 ----

    def _refill(self) -> None:
        """补充令牌（调用方需持有 _lock）。

        根据距上次补充的时间差计算新增令牌数，不超过 capacity。
        """
        now = time.monotonic()
        elapsed = now - self._last_refill
        if elapsed > 0 and self._rate > 0:
            new_tokens = elapsed * self._rate
            self._tokens = min(self._tokens + new_tokens, float(self._capacity))
        self._last_refill = now

    # ---- 公开方法 ----

    def consume(self, tokens: int = 1) -> bool:
        """尝试立即消费令牌（非阻塞）。

        Args:
            tokens: 要消费的令牌数，默认 1。

        Returns:
            True 消费成功；False 令牌不足。
        """
        if tokens <= 0:
            return True

        with self._lock:
            self._refill()
            if self._tokens >= tokens:
                self._tokens -= tokens
                return True
            return False

    def try_consume(self, tokens: int = 1, timeout: float = 2.0) -> bool:
        """尝试在指定时间内消费令牌（阻塞）。

        Args:
            tokens: 要消费的令牌数，默认 1。
            timeout: 最大等待秒数，默认 2.0。

        Returns:
            True 在超时前成功消费；False 超时。
        """
        if tokens <= 0:
            return True

        deadline = time.monotonic() + timeout
        while True:
            with self._lock:
                self._refill()
                if self._tokens >= tokens:
                    self._tokens -= tokens
                    return True

            # 释放锁后计算剩余等待时间
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return False
            # 等待到下一个令牌应该补充的时间点
            sleep_time = min(remaining, max(0.05, 1.0 / self._rate) if self._rate > 0 else remaining)
            time.sleep(sleep_time)

    def wait_for_token(self, timeout: float | None = None) -> bool:
        """等待并尝试消费 1 个令牌。

        Args:
            timeout: 最大等待秒数；None 表示无限等待（不推荐）。

        Returns:
            True 成功消费；False 超时。
        """
        return self.try_consume(tokens=1, timeout=timeout if timeout is not None else float("inf"))

    def get_status(self) -> dict[str, Any]:
        """获取令牌桶当前状态。

        Returns:
            包含 tokens, capacity, rate_per_minute, utilization 的字典。
        """
        with self._lock:
            self._refill()
            rate_per_minute = self._rate * 60.0
            utilization = 0.0 if self._capacity == 0 else (
                (self._capacity - self._tokens) / self._capacity
            )
            return {
                "tokens": round(self._tokens, 2),
                "capacity": self._capacity,
                "rate_per_minute": round(rate_per_minute, 1),
                "utilization": round(utilization, 4),
            }

    # ---- 属性 ----

    @property
    def rate(self) -> float:
        """当前令牌补充速率（令牌/秒）。"""
        return self._rate

    @property
    def capacity(self) -> int:
        """桶容量。"""
        return self._capacity