from __future__ import annotations
import time
from dataclasses import dataclass
from typing import Dict, Tuple, Optional

@dataclass
class Bucket:
    tokens: float
    last: float

class TokenBucketLimiter:
    """
    Per-key token bucket.
    Example: rate=6, per=3 => 6 actions per 3 seconds (bursty but capped).
    """
    def __init__(self, rate: float, per: float):
        self.rate = float(rate)
        self.per = float(per)
        self.capacity = float(rate)
        self._buckets: Dict[str, Bucket] = {}

    def _refill(self, b: Bucket, now: float):
        elapsed = now - b.last
        if elapsed <= 0:
            return
        b.last = now
        b.tokens = min(self.capacity, b.tokens + (elapsed * (self.rate / self.per)))

    def allow(self, key: str, cost: float = 1.0) -> Tuple[bool, float]:
        """
        Returns (allowed, retry_after_seconds)
        """
        now = time.monotonic()
        b = self._buckets.get(key)
        if b is None:
            b = Bucket(tokens=self.capacity, last=now)
            self._buckets[key] = b
        self._refill(b, now)
        if b.tokens >= cost:
            b.tokens -= cost
            return True, 0.0
        # compute retry
        missing = cost - b.tokens
        refill_rate = (self.rate / self.per)
        retry = missing / refill_rate if refill_rate > 0 else 1.0
        return False, max(0.2, retry)

class CooldownTable:
    """
    Simple cooldown: allow once per X seconds per key.
    """
    def __init__(self):
        self._next_ok: Dict[str, float] = {}

    def allow(self, key: str, seconds: float) -> Tuple[bool, float]:
        now = time.monotonic()
        nxt = self._next_ok.get(key, 0.0)
        if now >= nxt:
            self._next_ok[key] = now + seconds
            return True, 0.0
        return False, nxt - now

    def clear(self, key_prefix: Optional[str] = None):
        if key_prefix is None:
            self._next_ok.clear()
            return
        for k in list(self._next_ok.keys()):
            if k.startswith(key_prefix):
                self._next_ok.pop(k, None)