"""内存限流（MVP: 单进程 TokenBucket）。

未来切换 Redis: 实现 RedisRateLimiter，接口不变。
"""

import time
from collections import defaultdict

from app.core.config import get_settings


class _Bucket:
    __slots__ = ("tokens", "max_tokens", "refill_rate", "last_refill")

    def __init__(self, max_tokens: int, per_seconds: int) -> None:
        self.max_tokens = max_tokens
        self.refill_rate = max_tokens / per_seconds
        self.tokens = float(max_tokens)
        self.last_refill = time.monotonic()

    def allow(self) -> bool:
        now = time.monotonic()
        elapsed = now - self.last_refill
        self.tokens = min(self.max_tokens, self.tokens + elapsed * self.refill_rate)
        self.last_refill = now
        if self.tokens >= 1.0:
            self.tokens -= 1.0
            return True
        return False


# key → Bucket
_token_buckets: dict[str, _Bucket] = defaultdict(lambda: _Bucket(_token_limit, 60))
_ip_buckets: dict[str, _Bucket] = defaultdict(lambda: _Bucket(_ip_limit, 60))

# 延迟初始化
_token_limit = 60
_ip_limit = 30
_initialized = False


def _ensure_init() -> None:
    global _token_limit, _ip_limit, _initialized
    if _initialized:
        return
    settings = get_settings()
    _token_limit = settings.rate_limit_per_token_per_min
    _ip_limit = settings.rate_limit_per_ip_per_min
    # 重建 default dict factories
    _token_buckets.default_factory = lambda: _Bucket(_token_limit, 60)  # type: ignore[assignment]
    _ip_buckets.default_factory = lambda: _Bucket(_ip_limit, 60)  # type: ignore[assignment]
    _initialized = True


def check_rate_limit(token_key: str | None = None, ip_key: str | None = None) -> bool:
    """检查是否允许请求。返回 True 表示允许。"""
    _ensure_init()
    if token_key and not _token_buckets[token_key].allow():
        return False
    if ip_key and not _ip_buckets[ip_key].allow():
        return False
    return True