Trixy7/trixy_core/utils/ratelimit/decorators.py

# -*- coding: utf-8 -*-
"""
Rate-Limit-Decorators für Funktionen.

Bietet einfache Decorators für Durchsatzbegrenzung von
Funktionen und Methoden.
"""

from __future__ import annotations

import asyncio
import functools
import hashlib
from dataclasses import dataclass
from typing import Any, Callable, ParamSpec, TypeVar

from trixy_core.utils.ratelimit.token_bucket import TokenBucket, TokenBucketLimiter
from trixy_core.utils.ratelimit.sliding_window import SlidingWindow, SlidingWindowLimiter


P = ParamSpec("P")
R = TypeVar("R")


class RateLimitExceeded(Exception):
    """
    Exception wenn Rate-Limit überschritten wurde.

    Enthält Informationen zur Wartezeit bis zur nächsten
    erlaubten Ausführung.
    """

    def __init__(
        self,
        message: str = "Rate limit exceeded",
        retry_after: float | None = None,
        limit_key: str | None = None,
    ) -> None:
        """
        Initialisiert die Exception.

        Args:
            message: Fehlermeldung.
            retry_after: Zeit in Sekunden bis zur nächsten Erlaubnis.
            limit_key: Identifikator des Rate-Limits.
        """
        super().__init__(message)
        self.retry_after = retry_after
        self.limit_key = limit_key

    def __str__(self) -> str:
        """String-Repräsentation."""
        msg = super().__str__()
        if self.retry_after:
            msg += f" (retry after {self.retry_after:.2f}s)"
        return msg


@dataclass
class RateLimitConfig:
    """
    Konfiguration für Rate-Limiting.

    Definiert die Parameter für den Rate-Limiter.
    """

    # Token-Bucket-Parameter
    capacity: float = 10.0
    """Maximale Bucket-Kapazität (Token-Bucket)."""

    refill_rate: float = 1.0
    """Tokens pro Sekunde (Token-Bucket)."""

    # Sliding-Window-Parameter
    window_size: float = 60.0
    """Fenstergröße in Sekunden (Sliding-Window)."""

    max_requests: int = 60
    """Maximale Requests pro Fenster (Sliding-Window)."""

    # Allgemeine Parameter
    algorithm: str = "token_bucket"
    """Algorithmus: 'token_bucket' oder 'sliding_window'."""

    key_func: Callable[..., str] | None = None
    """Funktion zur Schlüsselgenerierung für Per-Key-Limiting."""

    block: bool = False
    """Ob gewartet werden soll statt Exception zu werfen."""

    max_wait: float | None = None
    """Maximale Wartezeit wenn block=True."""

    on_limited: Callable[[str], None] | None = None
    """Callback wenn Limit erreicht."""


# Globale Limiter-Instanzen
_token_bucket_limiters: dict[str, TokenBucketLimiter] = {}
_sliding_window_limiters: dict[str, SlidingWindowLimiter] = {}


def _get_key(
    func: Callable,
    args: tuple,
    kwargs: dict,
    key_func: Callable[..., str] | None,
) -> str:
    """Generiert einen Schlüssel für die Rate-Limit-Instanz."""
    if key_func:
        return key_func(*args, **kwargs)

    # Default: Funktionsname als Schlüssel
    return func.__qualname__


def rate_limit(
    capacity: float = 10.0,
    refill_rate: float = 1.0,
    window_size: float = 60.0,
    max_requests: int = 60,
    algorithm: str = "token_bucket",
    key_func: Callable[..., str] | None = None,
    block: bool = False,
    max_wait: float | None = None,
    on_limited: Callable[[str], None] | None = None,
    config: RateLimitConfig | None = None,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """
    Decorator für Rate-Limiting von Funktionen.

    Args:
        capacity: Token-Bucket-Kapazität.
        refill_rate: Token-Bucket Refill-Rate.
        window_size: Sliding-Window-Größe.
        max_requests: Sliding-Window Maximum.
        algorithm: 'token_bucket' oder 'sliding_window'.
        key_func: Optionale Funktion für Per-Key-Limiting.
        block: Ob gewartet werden soll.
        max_wait: Maximale Wartezeit.
        on_limited: Callback bei Limit-Erreichen.
        config: Alternative Konfiguration als Objekt.

    Returns:
        Decorator-Funktion.

    Example:
        @rate_limit(capacity=10, refill_rate=2)
        def api_call():
            return requests.get(url)

        @rate_limit(
            algorithm="sliding_window",
            window_size=60,
            max_requests=100,
            key_func=lambda user_id: str(user_id)
        )
        def user_action(user_id):
            pass

        @rate_limit(block=True, max_wait=5.0)
        def blocking_call():
            # Wartet bis zu 5 Sekunden
            pass
    """
    if config:
        cfg = config
    else:
        cfg = RateLimitConfig(
            capacity=capacity,
            refill_rate=refill_rate,
            window_size=window_size,
            max_requests=max_requests,
            algorithm=algorithm,
            key_func=key_func,
            block=block,
            max_wait=max_wait,
            on_limited=on_limited,
        )

    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        limiter_name = func.__qualname__

        # Limiter erstellen
        if cfg.algorithm == "token_bucket":
            if limiter_name not in _token_bucket_limiters:
                _token_bucket_limiters[limiter_name] = TokenBucketLimiter(
                    capacity=cfg.capacity,
                    refill_rate=cfg.refill_rate,
                )
            limiter = _token_bucket_limiters[limiter_name]
        else:
            if limiter_name not in _sliding_window_limiters:
                _sliding_window_limiters[limiter_name] = SlidingWindowLimiter(
                    window_size=cfg.window_size,
                    max_requests=cfg.max_requests,
                )
            limiter = _sliding_window_limiters[limiter_name]

        @functools.wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            key = _get_key(func, args, kwargs, cfg.key_func)

            if cfg.block:
                # Warten und dann ausführen
                if isinstance(limiter, TokenBucketLimiter):
                    success = limiter.wait_and_proceed(key, max_wait=cfg.max_wait)
                else:
                    success = limiter.wait_and_proceed(key, max_wait=cfg.max_wait)

                if not success:
                    if cfg.on_limited:
                        cfg.on_limited(key)
                    raise RateLimitExceeded(
                        "Rate limit exceeded after waiting",
                        limit_key=key,
                    )
            else:
                # Sofortige Prüfung
                if not limiter.is_allowed(key):
                    if cfg.on_limited:
                        cfg.on_limited(key)

                    # Retry-After berechnen
                    if isinstance(limiter, TokenBucketLimiter):
                        retry_after = limiter.get_wait_time(key)
                    else:
                        retry_after = limiter.get_retry_after(key)

                    raise RateLimitExceeded(
                        "Rate limit exceeded",
                        retry_after=retry_after,
                        limit_key=key,
                    )

            return func(*args, **kwargs)

        # Zugriff auf Limiter und Stats
        wrapper.limiter = limiter  # type: ignore
        wrapper.reset = lambda k=None: limiter.reset(k)  # type: ignore

        return wrapper

    return decorator


def async_rate_limit(
    capacity: float = 10.0,
    refill_rate: float = 1.0,
    window_size: float = 60.0,
    max_requests: int = 60,
    algorithm: str = "token_bucket",
    key_func: Callable[..., str] | None = None,
    block: bool = False,
    max_wait: float | None = None,
    on_limited: Callable[[str], None] | None = None,
    config: RateLimitConfig | None = None,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """
    Async-Decorator für Rate-Limiting.

    Wie @rate_limit, aber für async Funktionen.

    Args:
        Siehe rate_limit.

    Returns:
        Decorator-Funktion.

    Example:
        @async_rate_limit(capacity=10, refill_rate=2)
        async def api_call():
            return await aiohttp_session.get(url)
    """
    if config:
        cfg = config
    else:
        cfg = RateLimitConfig(
            capacity=capacity,
            refill_rate=refill_rate,
            window_size=window_size,
            max_requests=max_requests,
            algorithm=algorithm,
            key_func=key_func,
            block=block,
            max_wait=max_wait,
            on_limited=on_limited,
        )

    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        limiter_name = func.__qualname__

        if cfg.algorithm == "token_bucket":
            if limiter_name not in _token_bucket_limiters:
                _token_bucket_limiters[limiter_name] = TokenBucketLimiter(
                    capacity=cfg.capacity,
                    refill_rate=cfg.refill_rate,
                )
            limiter = _token_bucket_limiters[limiter_name]
        else:
            if limiter_name not in _sliding_window_limiters:
                _sliding_window_limiters[limiter_name] = SlidingWindowLimiter(
                    window_size=cfg.window_size,
                    max_requests=cfg.max_requests,
                )
            limiter = _sliding_window_limiters[limiter_name]

        @functools.wraps(func)
        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            key = _get_key(func, args, kwargs, cfg.key_func)

            if cfg.block:
                # Async warten
                start_wait = 0.0
                while True:
                    if limiter.is_allowed(key):
                        break

                    if isinstance(limiter, TokenBucketLimiter):
                        wait_time = limiter.get_wait_time(key)
                    else:
                        wait_time = limiter.get_retry_after(key)

                    if cfg.max_wait is not None:
                        if start_wait + wait_time > cfg.max_wait:
                            if cfg.on_limited:
                                cfg.on_limited(key)
                            raise RateLimitExceeded(
                                "Rate limit exceeded after waiting",
                                limit_key=key,
                            )

                    await asyncio.sleep(min(wait_time, 0.1))
                    start_wait += min(wait_time, 0.1)
            else:
                if not limiter.is_allowed(key):
                    if cfg.on_limited:
                        cfg.on_limited(key)

                    if isinstance(limiter, TokenBucketLimiter):
                        retry_after = limiter.get_wait_time(key)
                    else:
                        retry_after = limiter.get_retry_after(key)

                    raise RateLimitExceeded(
                        "Rate limit exceeded",
                        retry_after=retry_after,
                        limit_key=key,
                    )

            return await func(*args, **kwargs)

        wrapper.limiter = limiter  # type: ignore
        wrapper.reset = lambda k=None: limiter.reset(k)  # type: ignore

        return wrapper

    return decorator


def throttle(
    calls: int,
    period: float,
    key_func: Callable[..., str] | None = None,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """
    Vereinfachter Decorator für Throttling.

    Erlaubt eine bestimmte Anzahl von Aufrufen pro Zeitperiode.

    Args:
        calls: Maximale Anzahl Aufrufe.
        period: Zeitperiode in Sekunden.
        key_func: Optionale Funktion für Per-Key-Throttling.

    Returns:
        Decorator-Funktion.

    Example:
        @throttle(calls=10, period=60)  # 10 Aufrufe pro Minute
        def api_call():
            pass
    """
    return rate_limit(
        algorithm="sliding_window",
        window_size=period,
        max_requests=calls,
        key_func=key_func,
    )


def debounce(
    wait: float,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """
    Decorator für Debouncing.

    Verhindert mehrfache Aufrufe innerhalb einer kurzen Zeitspanne.
    Nur der letzte Aufruf innerhalb des Wartezeitraums wird ausgeführt.

    Args:
        wait: Wartezeit in Sekunden.

    Returns:
        Decorator-Funktion.

    Example:
        @debounce(0.5)
        def on_input_change(value):
            # Wird nur ausgeführt wenn 0.5s keine weiteren Aufrufe
            process(value)
    """
    import time
    import threading

    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        last_call = [0.0]
        timer: list[threading.Timer | None] = [None]
        lock = threading.Lock()
        last_args: list[tuple] = [()]
        last_kwargs: list[dict] = [{}]

        @functools.wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> None:
            with lock:
                last_args[0] = args
                last_kwargs[0] = kwargs

                if timer[0]:
                    timer[0].cancel()

                def execute():
                    func(*last_args[0], **last_kwargs[0])

                timer[0] = threading.Timer(wait, execute)
                timer[0].start()

        return wrapper  # type: ignore

    return decorator