from __future__ import annotations

import threading
import time
import unittest
from dataclasses import dataclass, field
from typing import Callable


Clock = Callable[[], float]
Sleeper = Callable[[float], None]


@dataclass
class TokenBucket:
    """Token-bucket rate limiter using a monotonic clock."""

    rate: float
    capacity: float
    clock: Clock = time.monotonic
    _tokens: float = field(init=False)
    _updated_at: float = field(init=False)
    _lock: threading.RLock = field(default_factory=threading.RLock, init=False)

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

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

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

    def consume(self, amount: float = 1.0) -> bool:
        """Try to consume tokens immediately."""
        self._validate_amount(amount)

        with self._lock:
            self._refill()
            if self._tokens < amount:
                return False

            self._tokens -= amount
            return True

    def time_until_available(self, amount: float = 1.0) -> float:
        """Return seconds until the requested token amount is available."""
        self._validate_amount(amount)

        with self._lock:
            self._refill()
            missing = amount - self._tokens
            return 0.0 if missing <= 0 else missing / self.rate

    def wait_and_consume(
        self,
        amount: float = 1.0,
        sleep: Sleeper = time.sleep,
    ) -> None:
        """Block until tokens can be consumed."""
        self._validate_amount(amount)

        while True:
            if self.consume(amount):
                return

            sleep(self.time_until_available(amount))

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

    def _validate_amount(self, amount: float) -> None:
        if amount <= 0:
            raise ValueError("amount must be positive")
        if amount > self.capacity:
            raise ValueError("amount cannot exceed bucket capacity")


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_consume_uses_available_tokens(self) -> None:
        clock = FakeClock()
        bucket = TokenBucket(rate=2.0, capacity=3.0, clock=clock)

        self.assertTrue(bucket.consume(3.0))
        self.assertFalse(bucket.consume())
        self.assertEqual(bucket.time_until_available(), 0.5)

    def test_tokens_refill_over_time_up_to_capacity(self) -> None:
        clock = FakeClock()
        bucket = TokenBucket(rate=2.0, capacity=5.0, clock=clock)

        self.assertTrue(bucket.consume(5.0))
        clock.advance(1.5)
        self.assertEqual(bucket.tokens, 3.0)

        clock.advance(10.0)
        self.assertEqual(bucket.tokens, 5.0)


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