"""
Rate limiting utilities for LLM API calls (Issue #39).

Provides token bucket rate limiter to prevent excessive API costs
and ensure compliance with provider rate limits.
"""

import asyncio
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional, Dict, Any
from collections import deque

logger = logging.getLogger(__name__)


def _utc_now() -> datetime:
    """Get current UTC time (timezone-aware)."""
    return datetime.now(timezone.utc)


# =============================================================================
# Cost Estimation
# =============================================================================

@dataclass
class ModelPricing:
    """Pricing information for a model (per 1M tokens)."""
    input_cost_per_1m: float  # Cost per 1M input tokens
    output_cost_per_1m: float  # Cost per 1M output tokens
    
    def estimate_cost(self, input_tokens: int, output_tokens: int) -> float:
        """Estimate cost for given token counts."""
        input_cost = (input_tokens / 1_000_000) * self.input_cost_per_1m
        output_cost = (output_tokens / 1_000_000) * self.output_cost_per_1m
        return input_cost + output_cost


# Bedrock pricing (as of 2024 - update as needed)
# https://aws.amazon.com/bedrock/pricing/
BEDROCK_PRICING: Dict[str, ModelPricing] = {
    # Llama 3.1 8B - Fast tier
    "llama-3.1-8b": ModelPricing(input_cost_per_1m=0.22, output_cost_per_1m=0.22),
    "llama-3.2-3b": ModelPricing(input_cost_per_1m=0.15, output_cost_per_1m=0.15),
    # Llama 3.3 70B - Smart tier  
    "llama-3.3-70b": ModelPricing(input_cost_per_1m=0.72, output_cost_per_1m=0.72),
    "llama-3.1-70b": ModelPricing(input_cost_per_1m=0.72, output_cost_per_1m=0.72),
    "llama-4-maverick": ModelPricing(input_cost_per_1m=0.50, output_cost_per_1m=1.50),
}

# OpenAI pricing
OPENAI_PRICING: Dict[str, ModelPricing] = {
    "gpt-4o": ModelPricing(input_cost_per_1m=2.50, output_cost_per_1m=10.00),
    "gpt-4o-mini": ModelPricing(input_cost_per_1m=0.15, output_cost_per_1m=0.60),
    "gpt-4-turbo": ModelPricing(input_cost_per_1m=10.00, output_cost_per_1m=30.00),
}

# Anthropic pricing
ANTHROPIC_PRICING: Dict[str, ModelPricing] = {
    "claude-sonnet-4-20250514": ModelPricing(input_cost_per_1m=3.00, output_cost_per_1m=15.00),
    "claude-3-5-sonnet-20241022": ModelPricing(input_cost_per_1m=3.00, output_cost_per_1m=15.00),
    "claude-3-haiku-20240307": ModelPricing(input_cost_per_1m=0.25, output_cost_per_1m=1.25),
}


def get_pricing(model: str, provider: str = "bedrock") -> Optional[ModelPricing]:
    """Get pricing for a model."""
    pricing_maps = {
        "bedrock": BEDROCK_PRICING,
        "openai": OPENAI_PRICING,
        "anthropic": ANTHROPIC_PRICING,
    }
    return pricing_maps.get(provider, {}).get(model)


# =============================================================================
# Token Bucket Rate Limiter
# =============================================================================

@dataclass
class RateLimitConfig:
    """Configuration for rate limiting."""
    requests_per_minute: int = 60  # RPM limit
    tokens_per_minute: int = 100_000  # TPM limit
    max_cost_per_hour: Optional[float] = None  # Cost limit (USD)
    enable_cost_logging: bool = True  # Log cost estimates
    warn_at_cost_percent: float = 0.8  # Warn when this % of hourly limit used


@dataclass 
class RateLimitStats:
    """Statistics from rate limiter."""
    total_requests: int = 0
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_estimated_cost: float = 0.0
    requests_delayed: int = 0
    total_delay_seconds: float = 0.0
    start_time: datetime = field(default_factory=_utc_now)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        runtime = (_utc_now() - self.start_time).total_seconds()
        return {
            "total_requests": self.total_requests,
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "total_tokens": self.total_input_tokens + self.total_output_tokens,
            "total_estimated_cost_usd": round(self.total_estimated_cost, 4),
            "requests_delayed": self.requests_delayed,
            "total_delay_seconds": round(self.total_delay_seconds, 2),
            "runtime_seconds": round(runtime, 2),
            "avg_requests_per_minute": round(self.total_requests / (runtime / 60), 2) if runtime > 0 else 0,
        }


class TokenBucketRateLimiter:
    """
    Token bucket rate limiter for LLM API calls.
    
    Implements both RPM (requests per minute) and TPM (tokens per minute)
    limits using the token bucket algorithm with smooth refilling.
    
    Example:
        limiter = TokenBucketRateLimiter(
            config=RateLimitConfig(
                requests_per_minute=60,
                tokens_per_minute=100_000,
                max_cost_per_hour=10.0
            ),
            model="llama-3.3-70b",
            provider="bedrock"
        )
        
        # Before each API call
        await limiter.acquire(estimated_tokens=1000)
        
        # After each API call
        limiter.record_usage(input_tokens=500, output_tokens=200)
    """
    
    def __init__(
        self,
        config: Optional[RateLimitConfig] = None,
        model: str = "llama-3.3-70b",
        provider: str = "bedrock",
    ):
        self.config = config or RateLimitConfig()
        self.model = model
        self.provider = provider
        self.pricing = get_pricing(model, provider)
        
        # Token buckets
        self._request_bucket = float(self.config.requests_per_minute)
        self._token_bucket = float(self.config.tokens_per_minute)
        self._last_refill = time.monotonic()
        
        # Cost tracking (sliding window for hourly limit)
        self._hourly_costs: deque = deque()  # (timestamp, cost) pairs
        
        # Statistics
        self.stats = RateLimitStats()
        
        # Lock for thread safety
        self._lock = asyncio.Lock()
    
    def _refill_buckets(self) -> None:
        """Refill token buckets based on elapsed time."""
        now = time.monotonic()
        elapsed = now - self._last_refill
        self._last_refill = now
        
        # Refill at rate of limit per minute
        request_refill = (elapsed / 60.0) * self.config.requests_per_minute
        token_refill = (elapsed / 60.0) * self.config.tokens_per_minute
        
        self._request_bucket = min(
            self._request_bucket + request_refill,
            float(self.config.requests_per_minute)
        )
        self._token_bucket = min(
            self._token_bucket + token_refill,
            float(self.config.tokens_per_minute)
        )
    
    def _get_hourly_cost(self) -> float:
        """Get total cost in the last hour."""
        now = time.time()
        hour_ago = now - 3600
        
        # Remove old entries
        while self._hourly_costs and self._hourly_costs[0][0] < hour_ago:
            self._hourly_costs.popleft()
        
        return sum(cost for _, cost in self._hourly_costs)
    
    def _check_cost_limit(self, estimated_cost: float) -> bool:
        """Check if we're within the hourly cost limit."""
        if self.config.max_cost_per_hour is None:
            return True
        
        current_hourly = self._get_hourly_cost()
        projected = current_hourly + estimated_cost
        
        # Warn if approaching limit
        if self.config.enable_cost_logging:
            usage_percent = current_hourly / self.config.max_cost_per_hour
            if usage_percent >= self.config.warn_at_cost_percent:
                logger.warning(
                    f"Rate limiter: {usage_percent:.1%} of hourly cost limit used "
                    f"(${current_hourly:.4f} / ${self.config.max_cost_per_hour:.2f})"
                )
        
        return projected <= self.config.max_cost_per_hour
    
    async def acquire(self, estimated_tokens: int = 1000) -> float:
        """
        Acquire permission to make an API call.
        
        Blocks until rate limits allow the request, or raises if cost limit exceeded.
        
        Args:
            estimated_tokens: Estimated total tokens for the request
            
        Returns:
            Delay in seconds that was waited (0 if no delay)
            
        Raises:
            RuntimeError: If hourly cost limit would be exceeded
        """
        async with self._lock:
            total_delay = 0.0
            
            while True:
                self._refill_buckets()
                
                # Check if we have capacity
                if self._request_bucket >= 1.0 and self._token_bucket >= estimated_tokens:
                    # Check cost limit
                    if self.pricing:
                        # Estimate cost (assume 50/50 input/output split for estimation)
                        estimated_cost = self.pricing.estimate_cost(
                            estimated_tokens // 2, 
                            estimated_tokens // 2
                        )
                        if not self._check_cost_limit(estimated_cost):
                            raise RuntimeError(
                                f"Hourly cost limit (${self.config.max_cost_per_hour:.2f}) "
                                f"would be exceeded. Current: ${self._get_hourly_cost():.4f}"
                            )
                    
                    # Consume from buckets
                    self._request_bucket -= 1.0
                    self._token_bucket -= estimated_tokens
                    
                    if total_delay > 0:
                        self.stats.requests_delayed += 1
                        self.stats.total_delay_seconds += total_delay
                        logger.debug(f"Rate limiter: delayed {total_delay:.2f}s")
                    
                    return total_delay
                
                # Calculate wait time
                wait_for_request = 0.0
                wait_for_tokens = 0.0
                
                if self._request_bucket < 1.0:
                    # Time to refill 1 request
                    wait_for_request = (1.0 - self._request_bucket) * 60.0 / self.config.requests_per_minute
                
                if self._token_bucket < estimated_tokens:
                    # Time to refill needed tokens
                    needed = estimated_tokens - self._token_bucket
                    wait_for_tokens = needed * 60.0 / self.config.tokens_per_minute
                
                wait_time = max(wait_for_request, wait_for_tokens, 0.1)  # Min 100ms
                wait_time = min(wait_time, 60.0)  # Max 60s
                
                logger.debug(
                    f"Rate limiter: waiting {wait_time:.2f}s "
                    f"(requests: {self._request_bucket:.1f}, tokens: {self._token_bucket:.0f})"
                )
                
                # Release lock while waiting
                self._lock.release()
                try:
                    await asyncio.sleep(wait_time)
                    total_delay += wait_time
                finally:
                    await self._lock.acquire()
    
    def record_usage(self, input_tokens: int, output_tokens: int) -> None:
        """
        Record actual token usage after an API call.
        
        Args:
            input_tokens: Actual input tokens used
            output_tokens: Actual output tokens used
        """
        self.stats.total_requests += 1
        self.stats.total_input_tokens += input_tokens
        self.stats.total_output_tokens += output_tokens
        
        # Calculate and record cost
        if self.pricing:
            cost = self.pricing.estimate_cost(input_tokens, output_tokens)
            self.stats.total_estimated_cost += cost
            self._hourly_costs.append((time.time(), cost))
            
            if self.config.enable_cost_logging:
                logger.info(
                    f"LLM call: {input_tokens} in + {output_tokens} out tokens, "
                    f"est. cost: ${cost:.6f}, total: ${self.stats.total_estimated_cost:.4f}"
                )
    
    def get_stats(self) -> Dict[str, Any]:
        """Get rate limiter statistics."""
        stats = self.stats.to_dict()
        stats["model"] = self.model
        stats["provider"] = self.provider
        stats["config"] = {
            "requests_per_minute": self.config.requests_per_minute,
            "tokens_per_minute": self.config.tokens_per_minute,
            "max_cost_per_hour": self.config.max_cost_per_hour,
        }
        stats["current_buckets"] = {
            "requests": round(self._request_bucket, 2),
            "tokens": round(self._token_bucket, 0),
        }
        return stats
    
    def reset_stats(self) -> None:
        """Reset statistics."""
        self.stats = RateLimitStats()
        self._hourly_costs.clear()


# =============================================================================
# Convenience Functions
# =============================================================================

def create_rate_limiter(
    model: str = "llama-3.3-70b",
    provider: str = "bedrock",
    rpm: int = 60,
    tpm: int = 100_000,
    max_cost_per_hour: Optional[float] = None,
    enable_cost_logging: bool = True,
) -> TokenBucketRateLimiter:
    """
    Create a rate limiter with the specified configuration.
    
    Args:
        model: Model name for cost estimation
        provider: Provider name ("bedrock", "openai", "anthropic")
        rpm: Requests per minute limit
        tpm: Tokens per minute limit
        max_cost_per_hour: Maximum cost per hour in USD (None = unlimited)
        enable_cost_logging: Log cost estimates for each call
        
    Returns:
        Configured TokenBucketRateLimiter
    """
    config = RateLimitConfig(
        requests_per_minute=rpm,
        tokens_per_minute=tpm,
        max_cost_per_hour=max_cost_per_hour,
        enable_cost_logging=enable_cost_logging,
    )
    return TokenBucketRateLimiter(config=config, model=model, provider=provider)
