import unittest
from collections.abc import Callable
from dataclasses import dataclass, field
from threading import Lock
from time import monotonic


@dataclass
class TokenBucket:
    capacity: float
    refill_rate: float
    clock: Callable[[], float] = monotonic
    _tokens: float = field(init=False)
    _updated_at: float = field(init=False)
    _lock: Lock = field(default_factory=Lock, init=False, repr=False)

    def __post_init__(self) -> None:
        if self.capacity <= 0:
            raise ValueError("capacity must be positive")
        if self.refill_rate <= 0:
            raise ValueError("refill_rate must be positive")

        self._tokens = self.capacity
        self._updated_at = self.clock()

    def _refill(self, now: float) -> None:
        # Monotonic clocks should not move backward, but custom clocks might.
        elapsed = max(0.0, now - self._updated_at)
        self._tokens = min(
            self.capacity,
            self._tokens + elapsed * self.refill_rate,
        )
        self._updated_at = max(self._updated_at, now)

    def consume(self, tokens: float = 1.0) -> bool:
        """Consume tokens if available, returning whether the request is allowed."""
        if tokens <= 0:
            raise ValueError("tokens must be positive")

        with self._lock:
            self._refill(self.clock())
            if tokens > self._tokens:
                return False

            self._tokens -= tokens
            return True

    def retry_after(self, tokens: float = 1.0) -> float:
        """Return the seconds until the requested number of tokens is available."""
        if tokens <= 0:
            raise ValueError("tokens must be positive")
        if tokens > self.capacity:
            return float("inf")

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


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

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

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


class TokenBucketTests(unittest.TestCase):
    def test_capacity_is_enforced(self) -> None:
        clock = FakeClock()
        bucket = TokenBucket(capacity=2, refill_rate=1, clock=clock)

        self.assertTrue(bucket.consume())
        self.assertTrue(bucket.consume())
        self.assertFalse(bucket.consume())

    def test_tokens_refill_over_time(self) -> None:
        clock = FakeClock()
        bucket = TokenBucket(capacity=2, refill_rate=0.5, clock=clock)

        self.assertTrue(bucket.consume(2))
        self.assertAlmostEqual(bucket.retry_after(), 2.0)

        clock.advance(2)
        self.assertTrue(bucket.consume())
        self.assertFalse(bucket.consume())


if __name__ == "__main__":
    unittest.main()