import asyncio
import time

class RateLimiter:
    """
    用于控制异步调用频率的类
    """
    def __init__(self, max_calls: int, period: float):
        """
        参数:
        - max_calls: 在指定周期内允许的最大调用次数。
        - period: 限制的周期，单位为秒。
        """
        self.max_calls = max_calls
        self.period = period
        self.semaphore = asyncio.Semaphore(self.max_calls)
        self.last_reset = time.time()

    async def __aenter__(self):
        await self.wait_for_slot()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self.release_slot()

    async def wait_for_slot(self):
        await self.semaphore.acquire()
        
        current_time = time.time()
        elapsed = current_time - self.last_reset
        
        # 如果一个周期已经过去，重置信号量
        if elapsed > self.period:
            # 释放所有已获取的槽位
            while self.semaphore._value < self.max_calls:
                self.semaphore.release()
            self.last_reset = current_time

    def release_slot(self):
        pass

class RateLimiter:
    """
    基于令牌桶算法的异步调用限速器，实现了平滑的请求频率控制。
    """
    def __init__(self, rate: float, capacity: int):
        """
        参数:
        - rate: 每秒生成的令牌数（即调用速率），例如 50次/60秒 -> 50/60次/秒
        - capacity: 令牌桶的容量，即可以瞬间处理的最大请求数。
        """
        self.rate = rate
        self.capacity = capacity
        
        # 当前可用的令牌数量
        self._tokens = capacity
        
        # 上次令牌更新时间
        self._last_update = time.time()
        
        self._lock = asyncio.Lock()

    async def __aenter__(self):
        await self.acquire()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass

    async def acquire(self):
        async with self._lock:
            now = time.time()
            
            # 计算自上次更新以来应该生成的新令牌
            elapsed = now - self._last_update
            new_tokens = elapsed * self.rate
            
            # 更新令牌数量，但不能超过容量
            self._tokens = min(self.capacity, self._tokens + new_tokens)
            self._last_update = now
            
            # 如果没有足够的令牌，需要等待
            if self._tokens < 1:
                # 计算需要等待的时间
                wait_time = (1 - self._tokens) / self.rate
                if wait_time > 0:
                    await asyncio.sleep(wait_time)
                    # 等待后更新令牌
                    now_after_wait = time.time()
                    elapsed_after_wait = now_after_wait - self._last_update
                    self._tokens = min(self.capacity, self._tokens + elapsed_after_wait * self.rate)
                    self._last_update = now_after_wait
            
            # 使用一个令牌
            self._tokens -= 1