"""Rate limiting with pluggable strategies.

Port of the Throttler/TokenBucketStrategy pattern from @coda/async.

Usage:
    throttler = Throttler(TokenBucketStrategy(capacity=1, refill_rate=2.0))
    throttler.acquire()  # blocks until a token is available
"""

import threading
import time
from typing import Protocol


class ThrottlerStrategy(Protocol):
    """Strategy interface for throttler behavior."""

    def try_acquire(self) -> tuple[bool, float]:
        """Attempt to acquire permission.

        Returns (success, retry_after_seconds). If success is True,
        retry_after_seconds is 0. If False, it indicates how long
        the caller should wait before retrying.
        """
        ...

    def clear(self) -> None:
        """Reset internal state."""
        ...


class TokenBucketStrategy:
    """Token bucket algorithm for rate limiting.

    Allows requests at a steady rate while supporting short bursts.
    Each acquisition consumes one token. Tokens refill continuously
    at refill_rate per second, up to capacity.

    Args:
        capacity: Maximum tokens the bucket can hold (burst size).
        refill_rate: Tokens added per second.
    """

    def __init__(self, capacity: int, refill_rate: float):
        if capacity < 1:
            raise ValueError('capacity must be >= 1')
        if refill_rate <= 0:
            raise ValueError('refill_rate must be > 0')

        self.capacity = capacity
        self.refill_rate = refill_rate
        self._tokens = float(capacity)
        self._refilled_at = time.monotonic()
        self._lock = threading.Lock()

    def try_acquire(self) -> tuple[bool, float]:
        with self._lock:
            now = time.monotonic()
            elapsed = now - self._refilled_at
            self._tokens = min(self.capacity, self._tokens + elapsed * self.refill_rate)
            self._refilled_at = now

            if self._tokens >= 1:
                self._tokens -= 1
                return True, 0.0

            deficit = 1.0 - self._tokens
            return False, deficit / self.refill_rate

    def clear(self) -> None:
        with self._lock:
            self._tokens = float(self.capacity)
            self._refilled_at = time.monotonic()


class Throttler:
    """Rate limiter that wraps a pluggable ThrottlerStrategy.

    Provides blocking acquire() and non-blocking try_acquire().
    """

    def __init__(self, strategy: ThrottlerStrategy):
        self._strategy = strategy

    def acquire(self) -> None:
        """Block until permission is granted."""
        while True:
            success, retry_after = self._strategy.try_acquire()
            if success:
                return
            time.sleep(retry_after)

    def try_acquire(self) -> bool:
        """Try to acquire permission without blocking."""
        success, _ = self._strategy.try_acquire()
        return success

    def clear(self) -> None:
        """Reset the underlying strategy."""
        self._strategy.clear()