src/api/rate_limiter.py

"""
Unified Rate Limiter

Provides configurable rate limiting for all API clients in the BeatDebate system.
Supports both per-second and per-hour rate limiting strategies.
"""

import asyncio
import time
from collections import deque
from typing import Optional, Union
from dataclasses import dataclass

import structlog

logger = structlog.get_logger(__name__)


@dataclass
class RateLimitConfig:
    """Configuration for rate limiting."""
    calls_per_second: Optional[float] = None
    calls_per_minute: Optional[int] = None
    calls_per_hour: Optional[int] = None
    burst_size: Optional[int] = None  # Maximum burst size for token bucket


class UnifiedRateLimiter:
    """
    Unified rate limiter supporting multiple rate limit strategies.
    
    Supports:
    - Per-second limiting (for LastFM: 3 requests/second)
    - Per-hour limiting (for Spotify: 50 requests/hour)
    - Token bucket algorithm for burst handling
    - Multiple time window tracking
    """
    
    def __init__(
        self,
        calls_per_second: Optional[float] = None,
        calls_per_minute: Optional[int] = None,
        calls_per_hour: Optional[int] = None,
        burst_size: Optional[int] = None,
        service_name: str = "api"
    ):
        """
        Initialize rate limiter with specified limits.
        
        Args:
            calls_per_second: Maximum calls per second (float for sub-second intervals)
            calls_per_minute: Maximum calls per minute
            calls_per_hour: Maximum calls per hour
            burst_size: Maximum burst size (defaults to calls_per_second * 2)
            service_name: Service name for logging
        """
        self.calls_per_second = calls_per_second
        self.calls_per_minute = calls_per_minute
        self.calls_per_hour = calls_per_hour
        self.service_name = service_name
        
        # Token bucket for burst handling
        if calls_per_second:
            self.burst_size = burst_size or max(int(calls_per_second * 2), 1)
            self.tokens = self.burst_size
            self.last_refill = time.time()
        else:
            self.burst_size = None
            self.tokens = 0
            self.last_refill = 0
        
        # Request history tracking
        self.request_times: deque = deque()
        self.lock = asyncio.Lock()
        
        self.logger = logger.bind(service=f"RateLimiter-{service_name}")
        
        self.logger.info(
            "Rate limiter initialized",
            calls_per_second=calls_per_second,
            calls_per_minute=calls_per_minute,
            calls_per_hour=calls_per_hour,
            burst_size=self.burst_size
        )
    
    @classmethod
    def for_lastfm(cls, calls_per_second: float = 3.0) -> "UnifiedRateLimiter":
        """
        Create rate limiter configured for Last.fm API.
        
        Args:
            calls_per_second: Calls per second (default: 3.0)
            
        Returns:
            Configured rate limiter for Last.fm
        """
        return cls(
            calls_per_second=calls_per_second,
            service_name="LastFM"
        )
    
    @classmethod
    def for_spotify(cls, calls_per_hour: int = 50) -> "UnifiedRateLimiter":
        """
        Create rate limiter configured for Spotify API.
        
        Args:
            calls_per_hour: Calls per hour (default: 50)
            
        Returns:
            Configured rate limiter for Spotify
        """
        return cls(
            calls_per_hour=calls_per_hour,
            service_name="Spotify"
        )
    
    @classmethod
    def for_gemini(cls, calls_per_minute: int = 15) -> "UnifiedRateLimiter":
        """
        Create rate limiter configured for Gemini LLM API.
        
        Args:
            calls_per_minute: Calls per minute (default: 15)
            
        Returns:
            Configured rate limiter for Gemini
        """
        return cls(
            calls_per_minute=calls_per_minute,
            service_name="Gemini"
        )
    
    async def wait_if_needed(self) -> None:
        """
        Wait if necessary to respect rate limits.
        
        This method should be called before making each API request.
        """
        async with self.lock:
            current_time = time.time()
            
            # Clean up old request times
            self._cleanup_old_requests(current_time)
            
            # Check each rate limit type
            wait_time = 0.0
            
            # Per-second rate limiting with token bucket
            if self.calls_per_second:
                wait_time = max(wait_time, await self._check_per_second_limit(current_time))
            
            # Per-minute rate limiting
            if self.calls_per_minute:
                wait_time = max(wait_time, self._check_per_minute_limit(current_time))
            
            # Per-hour rate limiting
            if self.calls_per_hour:
                wait_time = max(wait_time, self._check_per_hour_limit(current_time))
            
            # Wait if necessary
            if wait_time > 0:
                self.logger.debug(
                    "Rate limit wait required",
                    wait_time=wait_time,
                    current_requests=len(self.request_times)
                )
                await asyncio.sleep(wait_time)
                current_time = time.time()
            
            # Record this request
            self.request_times.append(current_time)
            
            # Consume token if using token bucket
            if self.calls_per_second and self.tokens > 0:
                self.tokens -= 1
    
    async def _check_per_second_limit(self, current_time: float) -> float:
        """
        Check per-second rate limit using token bucket algorithm.
        
        Args:
            current_time: Current timestamp
            
        Returns:
            Wait time in seconds (0 if no wait needed)
        """
        # Refill tokens based on time elapsed
        time_elapsed = current_time - self.last_refill
        tokens_to_add = time_elapsed * self.calls_per_second
        self.tokens = min(self.burst_size, self.tokens + tokens_to_add)
        self.last_refill = current_time
        
        # Check if we have tokens available
        if self.tokens >= 1:
            return 0.0
        
        # Calculate wait time for next token
        wait_time = (1.0 - self.tokens) / self.calls_per_second
        return wait_time
    
    def _check_per_minute_limit(self, current_time: float) -> float:
        """
        Check per-minute rate limit.
        
        Args:
            current_time: Current timestamp
            
        Returns:
            Wait time in seconds (0 if no wait needed)
        """
        minute_ago = current_time - 60
        recent_requests = [t for t in self.request_times if t > minute_ago]
        
        if len(recent_requests) < self.calls_per_minute:
            return 0.0
        
        # Find the oldest request in the current minute window
        oldest_request = min(recent_requests)
        wait_time = 60 - (current_time - oldest_request)
        return max(0.0, wait_time)
    
    def _check_per_hour_limit(self, current_time: float) -> float:
        """
        Check per-hour rate limit.
        
        Args:
            current_time: Current timestamp
            
        Returns:
            Wait time in seconds (0 if no wait needed)
        """
        hour_ago = current_time - 3600
        recent_requests = [t for t in self.request_times if t > hour_ago]
        
        if len(recent_requests) < self.calls_per_hour:
            return 0.0
        
        # Find the oldest request in the current hour window
        oldest_request = min(recent_requests)
        wait_time = 3600 - (current_time - oldest_request)
        return max(0.0, wait_time)
    
    def _cleanup_old_requests(self, current_time: float) -> None:
        """
        Remove old request timestamps that are outside all rate limit windows.
        
        Args:
            current_time: Current timestamp
        """
        # Determine the longest time window we need to track
        max_window = 0
        if self.calls_per_second:
            max_window = max(max_window, 60)  # Keep 1 minute for burst analysis
        if self.calls_per_minute:
            max_window = max(max_window, 60)
        if self.calls_per_hour:
            max_window = max(max_window, 3600)
        
        if max_window == 0:
            return
        
        # Remove requests older than the maximum window
        cutoff_time = current_time - max_window
        while self.request_times and self.request_times[0] < cutoff_time:
            self.request_times.popleft()
    
    def get_current_usage(self) -> dict:
        """
        Get current rate limit usage statistics.
        
        Returns:
            Dictionary with current usage information
        """
        current_time = time.time()
        
        # Count requests in each time window
        minute_ago = current_time - 60
        hour_ago = current_time - 3600
        
        minute_requests = len([t for t in self.request_times if t > minute_ago])
        hour_requests = len([t for t in self.request_times if t > hour_ago])
        
        usage = {
            "timestamp": current_time,
            "total_requests_tracked": len(self.request_times),
            "requests_last_minute": minute_requests,
            "requests_last_hour": hour_requests,
        }
        
        if self.calls_per_second:
            usage["tokens_available"] = self.tokens
            usage["burst_capacity"] = self.burst_size
            usage["calls_per_second_limit"] = self.calls_per_second
        
        if self.calls_per_minute:
            usage["calls_per_minute_limit"] = self.calls_per_minute
            usage["minute_usage_percent"] = (minute_requests / self.calls_per_minute) * 100
        
        if self.calls_per_hour:
            usage["calls_per_hour_limit"] = self.calls_per_hour
            usage["hour_usage_percent"] = (hour_requests / self.calls_per_hour) * 100
        
        return usage
    
    def reset(self) -> None:
        """Reset rate limiter state (useful for testing)."""
        self.request_times.clear()
        if self.calls_per_second:
            self.tokens = self.burst_size
            self.last_refill = time.time()
        
        self.logger.info("Rate limiter reset")


# Legacy compatibility - these will be deprecated
class RateLimiter(UnifiedRateLimiter):
    """Legacy rate limiter class for backward compatibility."""
    
    def __init__(self, calls_per_second: float):
        super().__init__(calls_per_second=calls_per_second, service_name="Legacy")


class SpotifyRateLimiter(UnifiedRateLimiter):
    """Legacy Spotify rate limiter class for backward compatibility."""
    
    def __init__(self, calls_per_hour: int):
        super().__init__(calls_per_hour=calls_per_hour, service_name="SpotifyLegacy")