"""
Rate Limiter - Token Bucket 算法

按 API 限制速度的 2 倍设置突发容量，按 API 限制速度持续补充。
"""
import time
import threading


class TokenBucket:
    def __init__(self, rpm: int, burst_factor: int = 2):
        self.capacity = rpm * burst_factor
        self.refill_rate = rpm / 60.0
        self.tokens = float(self.capacity)
        self.last_refill = time.monotonic()
        self._lock = threading.Lock()

    def acquire(self, tokens: int = 1, timeout: float = 300.0) -> bool:
        deadline = time.monotonic() + timeout
        while True:
            with self._lock:
                now = time.monotonic()
                elapsed = now - self.last_refill
                self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
                self.last_refill = now
                if self.tokens >= tokens:
                    self.tokens -= tokens
                    return True
                wait = (tokens - self.tokens) / self.refill_rate
            if time.monotonic() + wait > deadline:
                return False
            time.sleep(min(wait, 1.0))


class RateLimiter:
    """多 API 速率限制管理"""

    def __init__(self):
        self._buckets = {}

    def register(self, name: str, rpm: int, burst_factor: int = 2):
        self._buckets[name] = TokenBucket(rpm, burst_factor)

    def acquire(self, name: str, tokens: int = 1, timeout: float = 300.0) -> bool:
        bucket = self._buckets.get(name)
        if bucket is None:
            return True
        return bucket.acquire(tokens, timeout)