"""异步令牌桶限流器模块.

提供基于令牌桶算法的异步速率限制功能，
支持全局单例模式和异步上下文管理器。
"""

import asyncio
import time
from dataclasses import dataclass
from typing import Any, Optional


@dataclass
class RateLimitConfig:
    """限流配置数据类.

    Attributes:
        requests: 时间窗口内允许的最大请求数
        window_seconds: 时间窗口秒数
    """

    requests: int = 100
    window_seconds: int = 60


class TokenBucketRateLimiter:
    """异步令牌桶限流器.

    使用令牌桶算法实现速率限制，支持异步上下文管理器。
    全局单例模式确保整个应用共享同一限流器实例。

    Example:
        >>> limiter = TokenBucketRateLimiter()
        >>> async with limiter:
        ...     await make_api_call()

        >>> # 批量获取令牌
        >>> await limiter.acquire(tokens=5)
    """

    _instance: Optional["TokenBucketRateLimiter"] = None
    _lock: asyncio.Lock = asyncio.Lock()

    def __new__(cls, config: RateLimitConfig | None = None) -> "TokenBucketRateLimiter":
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, config: RateLimitConfig | None = None) -> None:
        """初始化限流器.

        Args:
            config: 限流配置，使用默认配置 if None
        """
        # 避免重复初始化
        if hasattr(self, "_initialized"):
            return

        self.config = config or RateLimitConfig()
        self.tokens: float = float(self.config.requests)
        self.last_update: float = time.monotonic()
        self._mutex: asyncio.Lock = asyncio.Lock()
        self._initialized: bool = True

    async def acquire(self, tokens: int = 1) -> None:
        """获取指定数量的令牌.

        如果令牌不足，会异步等待直到令牌补充足够。

        Args:
            tokens: 需要获取的令牌数量，默认为 1

        Raises:
            ValueError: tokens 小于等于 0
        """
        if tokens <= 0:
            raise ValueError("tokens must be positive")

        async with self._mutex:
            now = time.monotonic()
            elapsed = now - self.last_update

            # 补充令牌
            self.tokens = min(
                self.config.requests,
                self.tokens + elapsed * (self.config.requests / self.config.window_seconds),
            )
            self.last_update = now

            if self.tokens < tokens:
                # 计算需要等待的时间
                wait_time = (tokens - self.tokens) * (
                    self.config.window_seconds / self.config.requests
                )
                await asyncio.sleep(wait_time)
                self.tokens = 0
            else:
                self.tokens -= tokens

    async def __aenter__(self) -> "TokenBucketRateLimiter":
        """异步上下文管理器入口."""
        await self.acquire()
        return self

    async def __aexit__(
        self,
        exc_type: type | None,
        exc_val: BaseException | None,
        exc_tb: object | None,
    ) -> None:
        """异步上下文管理器出口."""
        pass

    def get_status(self) -> dict[str, Any]:
        """获取限流器当前状态.

        Returns:
            包含当前令牌数和配置信息的字典
        """
        return {
            "tokens": self.tokens,
            "max_tokens": self.config.requests,
            "window_seconds": self.config.window_seconds,
            "available": self.tokens / self.config.requests,
        }
