Trixy7/trixy_core/utils/ratelimit/token_bucket.py

# -*- coding: utf-8 -*-
"""
Token-Bucket-Algorithmus für Rate-Limiting.

Implementiert den klassischen Token-Bucket-Algorithmus für
gleichmäßige Durchsatzbegrenzung mit Burst-Unterstützung.
"""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any


@dataclass
class TokenBucketStats:
    """
    Statistiken für einen Token-Bucket.

    Bietet Einblick in die Nutzung und Auslastung des Buckets.
    """

    requests_allowed: int = 0
    """Anzahl erlaubter Requests."""

    requests_denied: int = 0
    """Anzahl abgelehnter Requests."""

    total_tokens_consumed: float = 0.0
    """Gesamtzahl verbrauchter Tokens."""

    last_request: datetime | None = None
    """Zeitpunkt des letzten Requests."""

    @property
    def total_requests(self) -> int:
        """Gesamtzahl der Requests."""
        return self.requests_allowed + self.requests_denied

    @property
    def allow_rate(self) -> float:
        """Erlaubnisrate in Prozent."""
        if self.total_requests == 0:
            return 100.0
        return (self.requests_allowed / self.total_requests) * 100

    def record_allow(self, tokens: float = 1.0) -> None:
        """Zeichnet einen erlaubten Request auf."""
        self.requests_allowed += 1
        self.total_tokens_consumed += tokens
        self.last_request = datetime.now()

    def record_deny(self) -> None:
        """Zeichnet einen abgelehnten Request auf."""
        self.requests_denied += 1
        self.last_request = datetime.now()

    def reset(self) -> None:
        """Setzt alle Statistiken zurück."""
        self.requests_allowed = 0
        self.requests_denied = 0
        self.total_tokens_consumed = 0.0
        self.last_request = None

    def to_dict(self) -> dict[str, Any]:
        """Konvertiert die Statistiken in ein Dictionary."""
        return {
            "requests_allowed": self.requests_allowed,
            "requests_denied": self.requests_denied,
            "total_requests": self.total_requests,
            "allow_rate": self.allow_rate,
            "total_tokens_consumed": self.total_tokens_consumed,
            "last_request": (
                self.last_request.isoformat()
                if self.last_request else None
            ),
        }


class TokenBucket:
    """
    Token-Bucket-Implementierung für Rate-Limiting.

    Der Bucket füllt sich kontinuierlich mit Tokens bis zur
    maximalen Kapazität. Jeder Request verbraucht Tokens.
    Wenn nicht genügend Tokens verfügbar sind, wird der
    Request abgelehnt.

    Attributes:
        capacity: Maximale Anzahl von Tokens im Bucket.
        refill_rate: Tokens pro Sekunde, die hinzugefügt werden.
        tokens: Aktuelle Anzahl von Tokens.

    Example:
        bucket = TokenBucket(capacity=10, refill_rate=2.0)

        if bucket.consume():
            # Request erlaubt
            process_request()
        else:
            # Rate limit erreicht
            raise TooManyRequests()
    """

    def __init__(
        self,
        capacity: float,
        refill_rate: float,
        initial_tokens: float | None = None,
    ) -> None:
        """
        Initialisiert den Token-Bucket.

        Args:
            capacity: Maximale Bucket-Kapazität.
            refill_rate: Tokens pro Sekunde.
            initial_tokens: Anfängliche Token-Anzahl (Default: capacity).
        """
        self._capacity = capacity
        self._refill_rate = refill_rate
        self._tokens = initial_tokens if initial_tokens is not None else capacity
        self._last_refill = time.monotonic()
        self._lock = threading.Lock()
        self._stats = TokenBucketStats()

    @property
    def capacity(self) -> float:
        """Gibt die Bucket-Kapazität zurück."""
        return self._capacity

    @property
    def refill_rate(self) -> float:
        """Gibt die Refill-Rate zurück."""
        return self._refill_rate

    @property
    def tokens(self) -> float:
        """Gibt die aktuelle Token-Anzahl zurück (nach Refill)."""
        with self._lock:
            self._refill()
            return self._tokens

    @property
    def stats(self) -> TokenBucketStats:
        """Gibt die Statistiken zurück."""
        return self._stats

    def _refill(self) -> None:
        """Füllt den Bucket basierend auf verstrichener Zeit."""
        now = time.monotonic()
        elapsed = now - self._last_refill

        if elapsed > 0:
            new_tokens = elapsed * self._refill_rate
            self._tokens = min(self._capacity, self._tokens + new_tokens)
            self._last_refill = now

    def consume(self, tokens: float = 1.0) -> bool:
        """
        Versucht, Tokens zu verbrauchen.

        Args:
            tokens: Anzahl zu verbrauchender Tokens.

        Returns:
            True wenn erfolgreich, False wenn nicht genug Tokens.
        """
        with self._lock:
            self._refill()

            if self._tokens >= tokens:
                self._tokens -= tokens
                self._stats.record_allow(tokens)
                return True

            self._stats.record_deny()
            return False

    def try_consume(self, tokens: float = 1.0) -> tuple[bool, float]:
        """
        Versucht, Tokens zu verbrauchen und gibt Wartezeit zurück.

        Args:
            tokens: Anzahl zu verbrauchender Tokens.

        Returns:
            Tuple (erfolgreich, Wartezeit in Sekunden).
        """
        with self._lock:
            self._refill()

            if self._tokens >= tokens:
                self._tokens -= tokens
                self._stats.record_allow(tokens)
                return True, 0.0

            # Berechne, wie lange auf genug Tokens gewartet werden muss
            needed = tokens - self._tokens
            wait_time = needed / self._refill_rate

            self._stats.record_deny()
            return False, wait_time

    def wait_and_consume(
        self,
        tokens: float = 1.0,
        max_wait: float | None = None,
    ) -> bool:
        """
        Wartet bis Tokens verfügbar sind und verbraucht sie.

        Args:
            tokens: Anzahl zu verbrauchender Tokens.
            max_wait: Maximale Wartezeit in Sekunden.

        Returns:
            True wenn erfolgreich, False wenn Timeout.
        """
        start = time.monotonic()

        while True:
            success, wait_time = self.try_consume(tokens)

            if success:
                return True

            # Timeout-Prüfung
            if max_wait is not None:
                elapsed = time.monotonic() - start
                remaining = max_wait - elapsed

                if remaining <= 0:
                    return False

                wait_time = min(wait_time, remaining)

            time.sleep(wait_time)

    def reset(self) -> None:
        """Setzt den Bucket auf volle Kapazität zurück."""
        with self._lock:
            self._tokens = self._capacity
            self._last_refill = time.monotonic()

    def drain(self) -> float:
        """
        Leert den Bucket und gibt verbleibende Tokens zurück.

        Returns:
            Anzahl der Tokens, die im Bucket waren.
        """
        with self._lock:
            self._refill()
            tokens = self._tokens
            self._tokens = 0
            return tokens

    def __repr__(self) -> str:
        """String-Repräsentation."""
        return (
            f"TokenBucket(capacity={self._capacity}, "
            f"refill_rate={self._refill_rate}, "
            f"tokens={self.tokens:.2f})"
        )


class TokenBucketLimiter:
    """
    Rate-Limiter basierend auf Token-Buckets.

    Verwaltet mehrere Buckets für verschiedene Schlüssel
    (z.B. User-IDs, IP-Adressen).
    """

    def __init__(
        self,
        capacity: float,
        refill_rate: float,
        key_ttl: float = 3600.0,
    ) -> None:
        """
        Initialisiert den Limiter.

        Args:
            capacity: Bucket-Kapazität pro Schlüssel.
            refill_rate: Tokens pro Sekunde pro Schlüssel.
            key_ttl: Zeit in Sekunden, bevor ungenutzte Buckets gelöscht werden.
        """
        self._capacity = capacity
        self._refill_rate = refill_rate
        self._key_ttl = key_ttl
        self._buckets: dict[str, tuple[TokenBucket, float]] = {}
        self._lock = threading.Lock()

    def _get_bucket(self, key: str) -> TokenBucket:
        """Holt oder erstellt einen Bucket für einen Schlüssel."""
        now = time.monotonic()

        with self._lock:
            if key in self._buckets:
                bucket, _ = self._buckets[key]
                self._buckets[key] = (bucket, now)
                return bucket

            bucket = TokenBucket(
                capacity=self._capacity,
                refill_rate=self._refill_rate,
            )
            self._buckets[key] = (bucket, now)
            return bucket

    def is_allowed(self, key: str, tokens: float = 1.0) -> bool:
        """
        Prüft, ob ein Request erlaubt ist.

        Args:
            key: Identifikator für den Rate-Limit-Kontext.
            tokens: Anzahl zu verbrauchender Tokens.

        Returns:
            True wenn erlaubt.
        """
        bucket = self._get_bucket(key)
        return bucket.consume(tokens)

    def get_wait_time(self, key: str, tokens: float = 1.0) -> float:
        """
        Gibt die Wartezeit bis zur Erlaubnis zurück.

        Args:
            key: Identifikator.
            tokens: Anzahl benötigter Tokens.

        Returns:
            Wartezeit in Sekunden (0 wenn sofort erlaubt).
        """
        bucket = self._get_bucket(key)
        success, wait_time = bucket.try_consume(tokens)
        return 0.0 if success else wait_time

    def wait_and_proceed(
        self,
        key: str,
        tokens: float = 1.0,
        max_wait: float | None = None,
    ) -> bool:
        """
        Wartet und verbraucht dann Tokens.

        Args:
            key: Identifikator.
            tokens: Anzahl zu verbrauchender Tokens.
            max_wait: Maximale Wartezeit.

        Returns:
            True wenn erfolgreich.
        """
        bucket = self._get_bucket(key)
        return bucket.wait_and_consume(tokens, max_wait)

    def cleanup(self) -> int:
        """
        Entfernt abgelaufene Buckets.

        Returns:
            Anzahl der entfernten Buckets.
        """
        now = time.monotonic()
        expired_keys = []

        with self._lock:
            for key, (_, last_access) in self._buckets.items():
                if now - last_access > self._key_ttl:
                    expired_keys.append(key)

            for key in expired_keys:
                del self._buckets[key]

        return len(expired_keys)

    def reset(self, key: str | None = None) -> None:
        """
        Setzt einen oder alle Buckets zurück.

        Args:
            key: Optionaler spezifischer Schlüssel.
        """
        with self._lock:
            if key:
                if key in self._buckets:
                    bucket, last_access = self._buckets[key]
                    bucket.reset()
            else:
                for bucket, _ in self._buckets.values():
                    bucket.reset()

    def get_stats(self, key: str) -> TokenBucketStats | None:
        """
        Gibt Statistiken für einen Schlüssel zurück.

        Args:
            key: Identifikator.

        Returns:
            Statistiken oder None.
        """
        with self._lock:
            if key in self._buckets:
                bucket, _ = self._buckets[key]
                return bucket.stats
            return None

    @property
    def bucket_count(self) -> int:
        """Gibt die Anzahl aktiver Buckets zurück."""
        with self._lock:
            return len(self._buckets)