"""Rate limiting via token bucket algorithm."""

import threading
import time
from typing import Dict, Any


class RateLimiter:
    """Token bucket rate limiter per agent/origin."""
    
    def __init__(self, tasa_tokens: float = 100.0, capacidad_max: float = 100.0):
        """
        Inicializa rate limiter con token bucket.
        
        Args:
            tasa_tokens: Tokens generados por segundo
            capacidad_max: Capacidad máxima del bucket
        """
        self.tasa_tokens = tasa_tokens
        self.capacidad_max = capacidad_max
        self.tokens = capacidad_max
        self.ultimo_relleno = time.time()
        self.lock = threading.Lock()
        self.limites_por_origen: Dict[str, Dict[str, float]] = {}
    
    def _rellenar_tokens(self):
        """Rellena tokens basado en tiempo transcurrido."""
        ahora = time.time()
        tiempo_transcurrido = ahora - self.ultimo_relleno
        tokens_a_añadir = tiempo_transcurrido * self.tasa_tokens
        self.tokens = min(self.capacidad_max, self.tokens + tokens_a_añadir)
        self.ultimo_relleno = ahora
    
    def permitir(self, costo: float = 1.0) -> bool:
        """
        Verifica si se puede consumir tokens.
        
        Args:
            costo: Tokens a consumir (default 1.0)
        
        Returns:
            True si se permite, False si está limitado
        """
        with self.lock:
            self._rellenar_tokens()
            if self.tokens >= costo:
                self.tokens -= costo
                return True
            return False
    
    def registrar_limite_origen(self, origen: str, tasa: float, capacidad: float):
        """Registra límite específico para un origen."""
        with self.lock:
            self.limites_por_origen[origen] = {
                "tasa": tasa,
                "capacidad": capacidad,
                "tokens": capacidad,
                "ultimo_relleno": time.time()
            }
    
    def permitir_por_origen(self, origen: str, costo: float = 1.0) -> bool:
        """Verifica límite para origen específico."""
        with self.lock:
            if origen not in self.limites_por_origen:
                # Si no tiene límite específico, usar global
                self._rellenar_tokens()
                if self.tokens >= costo:
                    self.tokens -= costo
                    return True
                return False
            
            limite = self.limites_por_origen[origen]
            ahora = time.time()
            tiempo_transcurrido = ahora - limite["ultimo_relleno"]
            tokens_a_añadir = tiempo_transcurrido * limite["tasa"]
            limite["tokens"] = min(limite["capacidad"], limite["tokens"] + tokens_a_añadir)
            limite["ultimo_relleno"] = ahora
            
            if limite["tokens"] >= costo:
                limite["tokens"] -= costo
                return True
            return False
    
    def obtener_estado(self) -> Dict[str, Any]:
        """Retorna estado actual del limiter."""
        with self.lock:
            return {
                "tokens_disponibles": self.tokens,
                "capacidad_max": self.capacidad_max,
                "tasa_tokens": self.tasa_tokens,
                "limites_por_origen": dict(self.limites_por_origen)
            }
