import time
from collections import defaultdict
from typing import Optional

class TokenBucket:
    """令牌桶限流器（每个 key 独立）"""
    def __init__(self, rate: float, burst: int):
        self.rate = rate          # 每秒令牌数
        self.burst = burst        # 最大令牌数
        self._tokens = defaultdict(float)
        self._last = defaultdict(float)

    def allow(self, key: str) -> bool:
        """返回 True 表示允许请求，False 表示限流"""
        now = time.time()
        # 补充令牌
        last = self._last.get(key, now)
        elapsed = now - last
        self._last[key] = now
        self._tokens[key] = min(self.burst, self._tokens.get(key, self.burst) + elapsed * self.rate)

        if self._tokens[key] >= 1.0:
            self._tokens[key] -= 1.0
            return True
        return False


class RateLimiter:
    """多 key 限流器，qps<=0 时放行"""
    def __init__(self, qps: int = 0, burst: int = 0):
        self.qps = qps
        self.burst = burst if burst > 0 else qps
        self._bucket = TokenBucket(rate=qps, burst=self.burst) if qps > 0 else None

    def allow(self, key: str) -> bool:
        if self._bucket is None:
            return True
        return self._bucket.allow(key)

    def active_buckets(self) -> int:
        """返回当前活跃的 key 数量（近似）"""
        if self._bucket is None:
            return 0
        return len(self._bucket._tokens)