"""
请求限流器
支持令牌桶、漏桶、滑动窗口算法
"""
import asyncio
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Optional, Callable, List, Any, Literal
from loguru import logger


@dataclass
class RateLimiterConfig:
    """限流器配置"""
    requests_per_second: float = 10.0
    algorithm: Literal["token-bucket", "leaky-bucket", "sliding-window"] = "token-bucket"
    bucket_size: Optional[int] = None  # 仅用于令牌桶


@dataclass
class RateLimiterStats:
    """限流器统计"""
    total_requests: int = 0
    total_wait_time: float = 0.0  # 毫秒
    waiting: int = 0
    available: int = 0


class RateLimiterBase(ABC):
    """限流器基类"""

    def __init__(self, requests_per_second: float):
        self.requests_per_second = requests_per_second

    @abstractmethod
    async def acquire(self) -> None:
        """获取请求许可"""
        pass

    @abstractmethod
    def get_stats(self) -> dict:
        """获取状态"""
        pass


class TokenBucketLimiter(RateLimiterBase):
    """令牌桶限流器"""

    def __init__(self, requests_per_second: float, bucket_size: Optional[int] = None):
        super().__init__(requests_per_second)
        self.bucket_size = bucket_size or int(requests_per_second)
        self.tokens = float(self.bucket_size)
        self.last_refill_time = time.time()
        self._waiting_queue: List[asyncio.Event] = []
        self._lock = asyncio.Lock()

    async def acquire(self) -> None:
        """获取请求许可"""
        async with self._lock:
            self._refill()

            if self.tokens >= 1:
                self.tokens -= 1
                return

        # 需要等待
        event = asyncio.Event()
        self._waiting_queue.append(event)

        # 计算等待时间
        time_to_next_token = (1 - self.tokens) / self.requests_per_second
        await asyncio.sleep(time_to_next_token)

        async with self._lock:
            self._refill()
            if self.tokens >= 1:
                self.tokens -= 1
                event.set()
                if event in self._waiting_queue:
                    self._waiting_queue.remove(event)

    def _refill(self) -> None:
        """补充令牌"""
        now = time.time()
        elapsed = now - self.last_refill_time
        tokens_to_add = elapsed * self.requests_per_second

        self.tokens = min(self.bucket_size, self.tokens + tokens_to_add)
        self.last_refill_time = now

    def get_stats(self) -> dict:
        """获取状态"""
        self._refill()
        return {
            "waiting": len(self._waiting_queue),
            "available": int(self.tokens),
        }


class LeakyBucketLimiter(RateLimiterBase):
    """漏桶限流器"""

    def __init__(self, requests_per_second: float):
        super().__init__(requests_per_second)
        self._queue: List[asyncio.Event] = []
        self._is_processing = False
        self._lock = asyncio.Lock()

    async def acquire(self) -> None:
        """获取请求许可"""
        event = asyncio.Event()

        async with self._lock:
            self._queue.append(event)
            if not self._is_processing:
                asyncio.create_task(self._process_queue())

        await event.wait()

    async def _process_queue(self) -> None:
        """处理队列"""
        self._is_processing = True
        interval = 1.0 / self.requests_per_second

        while True:
            async with self._lock:
                if not self._queue:
                    self._is_processing = False
                    return

                event = self._queue.pop(0)

            event.set()
            await asyncio.sleep(interval)

    def get_stats(self) -> dict:
        """获取状态"""
        return {
            "waiting": len(self._queue),
            "available": 1 if not self._queue else 0,
        }


class SlidingWindowLimiter(RateLimiterBase):
    """滑动窗口限流器"""

    def __init__(self, requests_per_second: float):
        super().__init__(requests_per_second)
        self.window_size = 1.0  # 1秒窗口
        self._requests: List[float] = []
        self._lock = asyncio.Lock()

    async def acquire(self) -> None:
        """获取请求许可"""
        async with self._lock:
            now = time.time()
            self._clean_old_requests(now)

            max_requests = int(self.requests_per_second) + 1

            if len(self._requests) < max_requests:
                self._requests.append(now)
                return

        # 需要等待
        oldest_request = self._requests[0]
        wait_time = oldest_request + self.window_size - time.time()

        if wait_time > 0:
            await asyncio.sleep(wait_time)

        async with self._lock:
            now = time.time()
            self._clean_old_requests(now)
            self._requests.append(now)

    def _clean_old_requests(self, now: float) -> None:
        """清理过期请求"""
        cutoff = now - self.window_size
        self._requests = [t for t in self._requests if t > cutoff]

    def get_stats(self) -> dict:
        """获取状态"""
        self._clean_old_requests(time.time())
        max_requests = int(self.requests_per_second) + 1
        return {
            "waiting": 0,
            "available": max(0, max_requests - len(self._requests)),
        }


class RateLimiter:
    """
    请求限流器
    支持多种限流算法
    """

    def __init__(self, config: Optional[RateLimiterConfig] = None):
        self.config = config or RateLimiterConfig()
        self._limiter = self._create_limiter()
        self._total_requests = 0
        self._total_wait_time = 0.0

    def _create_limiter(self) -> RateLimiterBase:
        """创建限流器实例"""
        algorithm = self.config.algorithm

        if algorithm == "leaky-bucket":
            return LeakyBucketLimiter(self.config.requests_per_second)
        elif algorithm == "sliding-window":
            return SlidingWindowLimiter(self.config.requests_per_second)
        else:  # token-bucket
            return TokenBucketLimiter(
                self.config.requests_per_second,
                self.config.bucket_size,
            )

    async def acquire(self) -> None:
        """获取请求许可"""
        start_time = time.time()
        await self._limiter.acquire()

        self._total_requests += 1
        self._total_wait_time += (time.time() - start_time) * 1000

    async def execute(self, fn: Callable) -> Any:
        """
        执行带限流的函数

        Args:
            fn: 要执行的函数

        Returns:
            Any: 函数返回值
        """
        await self.acquire()
        if asyncio.iscoroutinefunction(fn):
            return await fn()
        else:
            return fn()

    async def execute_batch(
        self,
        tasks: List[Callable],
        on_progress: Optional[Callable[[int, int], None]] = None,
    ) -> List[Any]:
        """
        批量执行带限流

        Args:
            tasks: 任务列表
            on_progress: 进度回调

        Returns:
            List[Any]: 结果列表
        """
        results = []

        for i, task in enumerate(tasks):
            result = await self.execute(task)
            results.append(result)

            if on_progress:
                if asyncio.iscoroutinefunction(on_progress):
                    await on_progress(i + 1, len(tasks))
                else:
                    on_progress(i + 1, len(tasks))

        return results

    def get_stats(self) -> RateLimiterStats:
        """获取统计信息"""
        current_stats = self._limiter.get_stats()
        return RateLimiterStats(
            total_requests=self._total_requests,
            total_wait_time=self._total_wait_time,
            waiting=current_stats["waiting"],
            available=current_stats["available"],
        )

    def reset_stats(self) -> None:
        """重置统计"""
        self._total_requests = 0
        self._total_wait_time = 0.0

    def update_rate(self, requests_per_second: float) -> None:
        """更新请求速率"""
        self.config.requests_per_second = requests_per_second
        self._limiter = self._create_limiter()
        logger.debug(f"限流器速率更新: {requests_per_second} req/s")


def create_rate_limiter(config: Optional[RateLimiterConfig] = None) -> RateLimiter:
    """创建限流器实例"""
    return RateLimiter(config)
