TokenBucket rate limiter with bursts

"""Token-bucket limiter: short bursts allowed, steady average rate enforced."""
import threading
import time


class TokenBucket:
    """Allow bursts up to capacity while refilling at rate tokens/second."""

    def __init__(self, rate: float, capacity: float = 1.0) -> None:
        if rate <= 0 or capacity <= 0:
            raise ValueError("rate and capacity must be positive")
        self.rate = rate
        self.capacity = float(capacity)
        self._tokens = self.capacity  # start full so the first burst is free
        self._updated = time.monotonic()
        self._lock = threading.Lock()

    def _refill(self) -> None:
        now = time.monotonic()
        elapsed = now - self._updated
        self._tokens = min(self.capacity, self._tokens + elapsed * self.rate)
        self._updated = now

    def try_acquire(self, tokens: int = 1) -> bool:
        """Consume tokens if available; never blocks, returns success."""
        with self._lock:
            self._refill()
            if self._tokens < tokens:
                return False
            self._tokens -= tokens
            return True

    def wait_time(self, tokens: int = 1) -> float:
        """Seconds until tokens fit, or 0.0 when they fit right now."""
        with self._lock:
            self._refill()
            return max(0.0, (tokens - self._tokens) / self.rate)