from __future__ import annotations

import asyncio
import time
import logging

logger = logging.getLogger("yuxi.channel.xmpp.rate_limiter")


class XmppRateLimiter:
    """Token Bucket 限流器，默认 5 msg/s，突发 10 条"""

    def __init__(self, max_rate: float = 5.0, burst: int = 10):
        self._rate = max_rate
        self._burst = burst
        self._tokens = float(burst)
        self._last_refill = time.monotonic()
        self._lock = asyncio.Lock()

    async def acquire(self) -> None:
        async with self._lock:
            now = time.monotonic()
            elapsed = now - self._last_refill
            self._tokens = min(self._burst, self._tokens + elapsed * self._rate)
            self._last_refill = now

            if self._tokens < 1.0:
                wait_time = (1.0 - self._tokens) / self._rate
                logger.debug("XMPP rate limit: waiting %.2fs", wait_time)
                await asyncio.sleep(wait_time)
                self._tokens = 0.0
            else:
                self._tokens -= 1.0


xmpp_rate_limiter = XmppRateLimiter(max_rate=5.0, burst=10)