from __future__ import annotations

import threading
import time
from collections.abc import Callable


class TokenBucket:
    """Thread-safe token-bucket rate limiter."""

    def __init__(
        self,
        rate: float,
        capacity: float,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        if rate <= 0:
            raise ValueError("rate must be positive")
        if capacity <= 0:
            raise ValueError("capacity must be positive")

        self.rate = rate
        self.capacity = capacity
        self._tokens = capacity
        self._clock = clock
        self._updated_at = clock()
        self._lock = threading.Lock()

    def _refill(self, now: float) -> None:
        elapsed = max(0.0, now - self._updated_at)
        self._tokens = min(self.capacity, self._tokens + elapsed * self.rate)
        self._updated_at = now

    def try_consume(self, tokens: float = 1.0) -> bool:
        """Consume tokens immediately, returning False when unavailable."""
        if tokens <= 0:
            raise ValueError("tokens must be positive")
        if tokens > self.capacity:
            return False

        with self._lock:
            self._refill(self._clock())
            if self._tokens < tokens:
                return False
            self._tokens -= tokens
            return True

    def wait_time(self, tokens: float = 1.0) -> float:
        """Return the seconds until the requested tokens are available."""
        if tokens <= 0:
            raise ValueError("tokens must be positive")
        if tokens > self.capacity:
            raise ValueError("requested tokens exceed bucket capacity")

        with self._lock:
            self._refill(self._clock())
            return max(0.0, (tokens - self._tokens) / self.rate)


def _run_tests() -> None:
    now = [0.0]
    bucket = TokenBucket(rate=2.0, capacity=3.0, clock=lambda: now[0])

    assert bucket.try_consume(3.0)
    assert not bucket.try_consume()
    assert bucket.wait_time() == 0.5

    # Two tokens should be restored after one second.
    now[0] = 1.0
    assert bucket.try_consume(2.0)
    assert not bucket.try_consume()


if __name__ == "__main__":
    _run_tests()