"""Rate and connection limits that protect the server from floods and abuse."""

from __future__ import annotations

import time
from collections import Counter
from collections.abc import Callable, Hashable

Clock = Callable[[], float]


class TokenBucket:
    """Allows ``rate`` actions per second on average with bursts up to ``burst``."""

    def __init__(self, rate: float, burst: float, clock: Clock = time.monotonic) -> None:
        self.rate = rate
        self.burst = burst
        self._clock = clock
        self._tokens = float(burst)
        self._updated = clock()

    def allow(self, cost: float = 1.0) -> bool:
        self._refill()
        if self._tokens >= cost:
            self._tokens -= cost
            return True
        return False

    @property
    def full(self) -> bool:
        self._refill()
        return self._tokens >= self.burst

    def _refill(self) -> None:
        now = self._clock()
        self._tokens = min(self.burst, self._tokens + (now - self._updated) * self.rate)
        self._updated = now


class KeyedRateLimiter:
    """One token bucket per key (e.g. per client IP)."""

    def __init__(self, rate: float, burst: float, clock: Clock = time.monotonic) -> None:
        self.rate = rate
        self.burst = burst
        self._clock = clock
        self._buckets: dict[Hashable, TokenBucket] = {}

    def allow(self, key: Hashable) -> bool:
        bucket = self._buckets.get(key)
        if bucket is None:
            bucket = self._buckets[key] = TokenBucket(self.rate, self.burst, self._clock)
        return bucket.allow()

    def prune(self) -> None:
        """Forget keys whose bucket has refilled completely (they behave like new keys)."""
        for key in [key for key, bucket in self._buckets.items() if bucket.full]:
            del self._buckets[key]

    def __len__(self) -> int:
        return len(self._buckets)


class ConnectionLimiter:
    """Counts open connections in total and per client IP."""

    def __init__(self, max_total: int, max_per_ip: int) -> None:
        self.max_total = max_total
        self.max_per_ip = max_per_ip
        self._per_ip: Counter[str] = Counter()
        self._total = 0

    @property
    def total(self) -> int:
        return self._total

    def acquire(self, ip: str) -> bool:
        if self._total >= self.max_total or self._per_ip[ip] >= self.max_per_ip:
            return False
        self._per_ip[ip] += 1
        self._total += 1
        return True

    def release(self, ip: str) -> None:
        if self._per_ip[ip] <= 0:
            return
        self._per_ip[ip] -= 1
        self._total -= 1
        if self._per_ip[ip] == 0:
            del self._per_ip[ip]