import asyncio
import time


class RateLimiter:
    """Async token bucket rate limiter.

    Allows a burst of requests, then throttles to a steady rate.
    Thread-safe via asyncio.Lock.
    """

    def __init__(self, rate: float = 2.0, burst: int = 3):
        """
        Args:
            rate: Minimum seconds between requests (after burst exhausted).
            burst: Maximum number of requests allowed immediately.
        """
        self._rate = rate
        self._burst = burst
        self._tokens = float(burst)
        self._last_refill = time.monotonic()
        self._lock: asyncio.Lock | None = None

    def _get_lock(self) -> asyncio.Lock:
        if self._lock is None:
            self._lock = asyncio.Lock()
        return self._lock

    async def acquire(self) -> None:
        """Acquire a token, waiting if necessary."""
        while True:
            async with self._get_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:
                    self._tokens -= 1.0
                    return

            await asyncio.sleep(self._rate * 0.5)

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

    async def __aexit__(self, *exc):
        pass
