"""Token-bucket rate limiter."""
from __future__ import annotations

import asyncio
import time
import logging

logger = logging.getLogger("runblob.rate_limiter")


class TokenBucketRateLimiter:
    """Async token-bucket rate limiter."""

    __slots__ = (
        "_rate",
        "_max_tokens",
        "_tokens",
        "_last_refill",
        "_lock",
    )

    def __init__(self, rps: float = 50.0, burst: float | None = None) -> None:
        """
        Args:
            rps: Sustained requests per second.
            burst: Maximum burst size. Defaults to 2× rps.
        """
        self._rate = rps
        self._max_tokens = burst if burst is not None else rps * 2
        self._tokens = self._max_tokens
        self._last_refill = time.monotonic()
        self._lock = asyncio.Lock()

    def _refill(self) -> None:
        """Add tokens based on elapsed time since last refill."""
        now = time.monotonic()
        elapsed = now - self._last_refill
        new_tokens = elapsed * self._rate
        if new_tokens > 0:
            self._tokens = min(self._max_tokens, self._tokens + new_tokens)
            self._last_refill = now

    async def acquire(self, tokens: float = 1.0) -> None:
        """Acquire tokens, sleeping if necessary until available."""
        while True:
            async with self._lock:
                self._refill()
                if self._tokens >= tokens:
                    self._tokens -= tokens
                    return
                # Calculate wait time until enough tokens are available
                deficit = tokens - self._tokens
                wait_time = deficit / self._rate

            await asyncio.sleep(wait_time)

    @property
    def available_tokens(self) -> float:
        """Current token count (approximate, not locked)."""
        self._refill()
        return self._tokens
