from __future__ import annotations

import time
from collections.abc import Callable


class TokenBucket:
    """Token-bucket rate limiter.

    `capacity` is the burst size. `refill_rate` is tokens added per second.
    """

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

        self.capacity = capacity
        self.refill_rate = refill_rate
        self._clock = clock
        self._sleep = sleep
        self._tokens = capacity
        self._updated_at = clock()

    @property
    def tokens(self) -> float:
        self._refill()
        return self._tokens

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

    def consume(self, tokens: float = 1.0) -> bool:
        """Try to spend tokens immediately."""
        if tokens <= 0:
            raise ValueError("tokens must be positive")
        if tokens > self.capacity:
            raise ValueError("tokens cannot exceed bucket capacity")

        self._refill()
        if self._tokens < tokens:
            return False

        self._tokens -= tokens
        return True

    def wait_time(self, tokens: float = 1.0) -> float:
        """Return seconds until `tokens` can be consumed."""
        if tokens <= 0:
            raise ValueError("tokens must be positive")
        if tokens > self.capacity:
            raise ValueError("tokens cannot exceed bucket capacity")

        self._refill()
        missing = max(0.0, tokens - self._tokens)
        return missing / self.refill_rate

    def acquire(self, tokens: float = 1.0) -> None:
        """Block until tokens are available, then consume them."""
        while not self.consume(tokens):
            self._sleep(self.wait_time(tokens))


class FakeClock:
    def __init__(self) -> None:
        self.current = 0.0

    def __call__(self) -> float:
        return self.current

    def sleep(self, seconds: float) -> None:
        self.current += seconds

    def advance(self, seconds: float) -> None:
        self.current += seconds


def test_consume_respects_capacity() -> None:
    clock = FakeClock()
    bucket = TokenBucket(capacity=2.0, refill_rate=1.0, clock=clock, sleep=clock.sleep)

    assert bucket.consume()
    assert bucket.consume()
    assert not bucket.consume()
    assert bucket.tokens == 0.0


def test_refill_and_acquire() -> None:
    clock = FakeClock()
    bucket = TokenBucket(capacity=3.0, refill_rate=2.0, clock=clock, sleep=clock.sleep)

    assert bucket.consume(3.0)
    assert bucket.wait_time(1.0) == 0.5

    clock.advance(0.25)
    assert not bucket.consume(1.0)

    bucket.acquire(1.0)
    assert clock.current == 0.5
    assert bucket.tokens == 0.0


def run_tests() -> None:
    test_consume_respects_capacity()
    test_refill_and_acquire()
    print("tests passed")


if __name__ == "__main__":
    run_tests()