"""Thread-safe token bucket rate limiter.

Correctness notes:
- Refill is computed from elapsed wall-clock time using a monotonic clock.
- Tokens are capped at capacity.
- acquire(n) blocks until n tokens available or timeout; partial consumes never happen.
- A single lock guards state; condition variable for waiters.
"""
from __future__ import annotations

import threading
import time
from typing import Optional


class TokenBucket:
    def __init__(self, rate: float, capacity: float):
        if rate <= 0:
            raise ValueError("rate must be > 0")
        if capacity <= 0:
            raise ValueError("capacity must be > 0")
        if rate > capacity:
            # allow; but typically rate <= capacity. No restriction.
            pass
        self._rate = float(rate)
        self._capacity = float(capacity)
        self._tokens = float(capacity)
        self._last = time.monotonic()
        self._lock = threading.Lock()
        self._cond = threading.Condition(self._lock)

    @property
    def rate(self) -> float:
        return self._rate

    @property
    def capacity(self) -> float:
        return self._capacity

    def _refill(self, now: float) -> None:
        elapsed = now - self._last
        if elapsed <= 0:
            return
        self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
        self._last = now

    def tokens(self) -> float:
        with self._lock:
            self._refill(time.monotonic())
            return self._tokens

    def try_acquire(self, n: float = 1.0) -> bool:
        if n <= 0:
            raise ValueError("n must be > 0")
        if n > self._capacity:
            raise ValueError("n exceeds capacity")
        with self._lock:
            now = time.monotonic()
            self._refill(now)
            if self._tokens >= n - 1e-12:
                self._tokens -= n
                return True
            return False

    def acquire(self, n: float = 1.0, timeout: Optional[float] = None) -> bool:
        if n <= 0:
            raise ValueError("n must be > 0")
        if n > self._capacity:
            raise ValueError("n exceeds capacity")
        deadline = None if timeout is None else time.monotonic() + timeout
        with self._lock:
            while True:
                now = time.monotonic()
                self._refill(now)
                if self._tokens >= n - 1e-12:
                    self._tokens -= n
                    return True
                # time until enough tokens
                deficit = n - self._tokens
                wait = deficit / self._rate
                if deadline is not None:
                    remaining = deadline - now
                    if remaining <= 0:
                        return False
                    wait = min(wait, remaining)
                # Guard against sub-millisecond waits that degenerate into a
                # busy spin on spurious wakeups (cond.wait(1e-15) returns at
                # once). Never wait longer than the allowed remaining time.
                if wait < 1e-3:
                    if deadline is not None:
                        if remaining <= 1e-3:
                            return False
                        wait = min(remaining, 1e-3)
                    else:
                        wait = 1e-3
                self._cond.wait(timeout=wait)
                # loop rechecks

    def reset(self, tokens: Optional[float] = None) -> None:
        with self._lock:
            self._tokens = self._capacity if tokens is None else float(tokens)
            self._last = time.monotonic()
            self._cond.notify_all()