"""Per-source rate limits — a token bucket per source, settings from config/sources.yaml.

A caller that would have to wait longer than its budget does not wait: it gets ``False`` and the
client turns that into a ``rate_limited`` source error. A limit must never become a retry loop
that outlasts the orchestrator's call limit (ADR-0001 §8).

Tokens are reserved under the lock and slept for outside it, so concurrent callers queue in
order without holding the lock while they sleep.
"""

from __future__ import annotations

import threading
import time
from collections.abc import Callable

from src import config


class TokenBucket:
    def __init__(
        self,
        rate_per_second: float,
        burst: int,
        *,
        clock: Callable[[], float] = time.monotonic,
        sleep: Callable[[float], None] = time.sleep,
    ):
        if rate_per_second <= 0 or burst < 1:
            raise ValueError("rate must be positive and burst at least 1")
        self.rate = float(rate_per_second)
        self.burst = int(burst)
        self._clock = clock
        self._sleep = sleep
        self._tokens = float(burst)
        self._last = clock()
        self._lock = threading.Lock()

    def acquire(self, max_wait: float) -> bool:
        """Take one token, waiting at most ``max_wait`` seconds. False if that is not enough."""
        with self._lock:
            now = self._clock()
            self._tokens = min(self.burst, self._tokens + (now - self._last) * self.rate)
            self._last = now
            wait = 0.0 if self._tokens >= 1 else (1 - self._tokens) / self.rate
            if wait > max_wait:
                return False
            self._tokens -= 1  # reserve now; may go negative while the caller sleeps
        if wait > 0:
            self._sleep(wait)
        return True


class RateLimiters:
    """One bucket per declared source, built lazily from config."""

    def __init__(
        self,
        *,
        clock: Callable[[], float] = time.monotonic,
        sleep: Callable[[float], None] = time.sleep,
    ):
        self._clock = clock
        self._sleep = sleep
        self._buckets: dict[str, TokenBucket] = {}
        self._lock = threading.Lock()

    def bucket(self, source: str) -> TokenBucket:
        with self._lock:
            if source not in self._buckets:
                cfg = config.source(source)
                self._buckets[source] = TokenBucket(
                    cfg["rate_per_second"], cfg["burst"], clock=self._clock, sleep=self._sleep
                )
            return self._buckets[source]

    def acquire(self, source: str, max_wait: float) -> bool:
        return self.bucket(source).acquire(max_wait)