"""API 速率限制器实现。

提供基于固定时间间隔的速率限制，强制两次请求之间保持最小时间间隔。
适合 Tushare 等需要严格控制请求频率的 API。
"""

import time
import threading
from typing import Optional
from dataclasses import dataclass, field


@dataclass
class RateLimiterStats:
    """速率限制器统计信息。"""

    total_requests: int = 0
    successful_requests: int = 0
    denied_requests: int = 0
    total_wait_time: float = 0.0
    last_request_time: Optional[float] = None  # 上次请求开始时间


class TokenBucketRateLimiter:
    """基于固定时间间隔的速率限制器。

    强制两次请求之间保持最小时间间隔，无论请求处理耗时多久。
    适合需要严格控制请求频率、避免触发服务端限流的场景。

    Attributes:
        rate_limit: 每分钟允许的请求数
        min_interval: 两次请求之间的最小时间间隔（秒）
    """

    def __init__(
        self,
        rate_limit: int = 150,
        **kwargs,
    ) -> None:
        """初始化速率限制器。

        Args:
            rate_limit: 每分钟允许的请求数（默认 150）
        """
        self.rate_limit = rate_limit
        # 计算最小间隔：60秒 / 每分钟请求数
        self.min_interval = 60.0 / rate_limit

        self._lock = threading.RLock()
        self._stats = RateLimiterStats()

    def acquire(self, timeout: float = float("inf")) -> tuple[bool, float]:
        """获取请求许可，确保与上次请求间隔足够时间。

        会等待直到距离上次请求的时间 >= min_interval。
        注意：全程加锁，确保多线程下严格串行执行。

        Args:
            timeout: 最大等待时间（秒），默认无限等待

        Returns:
            (success, wait_time): 是否成功，以及实际等待时间
        """
        start_time = time.monotonic()

        with self._lock:
            now = time.monotonic()

            # 计算距离上次请求的时间
            if self._stats.last_request_time is not None:
                elapsed = now - self._stats.last_request_time
                time_to_wait = self.min_interval - elapsed

                if time_to_wait > 0:
                    # 需要等待
                    if timeout != float("inf") and time_to_wait > timeout:
                        # 超过最大等待时间
                        self._stats.total_requests += 1
                        self._stats.denied_requests += 1
                        return False, time_to_wait

                    # 在锁内等待（全程加锁，确保多线程严格串行）
                    time.sleep(time_to_wait)
                    now = time.monotonic()

            # 更新上次请求时间（请求开始前）
            self._stats.last_request_time = now
            wait_time = now - start_time

            self._stats.total_requests += 1
            self._stats.successful_requests += 1
            self._stats.total_wait_time += wait_time

            return True, wait_time

    def acquire_nonblocking(self) -> tuple[bool, float]:
        """尝试非阻塞地获取请求许可。

        Returns:
            (success, wait_time): 是否成功，以及需要等待的时间
        """
        with self._lock:
            now = time.monotonic()

            if self._stats.last_request_time is not None:
                elapsed = now - self._stats.last_request_time
                time_to_wait = self.min_interval - elapsed

                if time_to_wait > 0:
                    self._stats.total_requests += 1
                    self._stats.denied_requests += 1
                    return False, time_to_wait

            # 立即获得许可
            self._stats.last_request_time = now
            self._stats.total_requests += 1
            self._stats.successful_requests += 1
            return True, 0.0

    def get_min_interval(self) -> float:
        """获取最小请求间隔。

        Returns:
            两次请求之间的最小时间间隔（秒）
        """
        return self.min_interval

    def get_time_until_next_request(self) -> float:
        """获取距离下次允许请求的时间。

        Returns:
            距离下次请求还需要等待的时间（秒），0 表示可以立即请求
        """
        with self._lock:
            if self._stats.last_request_time is None:
                return 0.0

            elapsed = time.monotonic() - self._stats.last_request_time
            return max(0.0, self.min_interval - elapsed)

    def get_stats(self) -> RateLimiterStats:
        """获取速率限制器统计信息。

        Returns:
            RateLimiterStats 实例的副本
        """
        with self._lock:
            return RateLimiterStats(
                total_requests=self._stats.total_requests,
                successful_requests=self._stats.successful_requests,
                denied_requests=self._stats.denied_requests,
                total_wait_time=self._stats.total_wait_time,
                last_request_time=self._stats.last_request_time,
            )

    def reset(self) -> None:
        """重置限流器状态（调试用）。"""
        with self._lock:
            self._stats = RateLimiterStats()
