"""
LLM management module with advanced features.

Provides:
- Model registration and management
- Retry mechanisms with exponential backoff
- Circuit breaker pattern
- Rate limiting
- Token usage tracking
- Response caching
- Model fallback chain
"""

from typing import Dict, Any, Callable, Optional, List
import time
from datetime import datetime
from collections import defaultdict
import hashlib


class CircuitBreaker:
    """Circuit breaker to prevent cascading failures."""
    
    def __init__(self, 
                 failure_threshold: int = 5,
                 recovery_timeout: int = 60):
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.failure_count = 0
        self.last_failure_time: Optional[float] = None
        self.state = "closed"  # closed, open, half-open
    
    def call(self, func: Callable, *args, **kwargs) -> Any:
        """Execute function with circuit breaker protection."""
        if self.state == "open":
            # Check if recovery timeout has passed
            if (self.last_failure_time and 
                time.time() - self.last_failure_time > self.recovery_timeout):
                self.state = "half-open"
                self.failure_count = 0
            else:
                raise Exception("Circuit breaker is OPEN")
        
        try:
            result = func(*args, **kwargs)
            
            # Success - reset if in half-open state
            if self.state == "half-open":
                self.state = "closed"
                self.failure_count = 0
            
            return result
            
        except Exception as e:
            self.failure_count += 1
            self.last_failure_time = time.time()
            
            if self.failure_count >= self.failure_threshold:
                self.state = "open"
            
            raise e


class LLMManager:
    """
    Enhanced LLM Manager with reliability and monitoring features.
    
    Features:
    - Model registry with metadata
    - Retry with exponential backoff
    - Circuit breaker pattern
    - Response caching
    - Token usage tracking
    - Fallback chain
    - Performance metrics
    """
    
    def __init__(self, 
                 max_retries: int = 3,
                 enable_caching: bool = True):
        self.models: Dict[str, Callable[[str, Dict[str, Any]], str]] = {}
        self.model_metadata: Dict[str, Dict[str, Any]] = {}
        self.active_model: Optional[str] = None
        self.fallback_chain: List[str] = []
        
        # Features
        self.max_retries = max_retries
        self.enable_caching = enable_caching
        
        # Circuit breakers per model
        self.circuit_breakers: Dict[str, CircuitBreaker] = {}
        
        # Response cache
        self.cache: Dict[str, Any] = {}
        
        # Metrics
        self.metrics = {
            'total_requests': 0,
            'successful_requests': 0,
            'failed_requests': 0,
            'cache_hits': 0,
            'total_retries': 0,
            'total_tokens': 0
        }
        
        self.model_stats: Dict[str, Dict[str, Any]] = defaultdict(lambda: {
            'requests': 0,
            'successes': 0,
            'failures': 0,
            'total_latency': 0.0,
            'avg_latency': 0.0
        })

    def register_model(self, 
                      name: str, 
                      inference_fn: Callable[[str, Dict[str, Any]], str],
                      metadata: Dict[str, Any] = None):
        """
        Register an LLM model.
        
        Args:
            name: Model name
            inference_fn: Function to call for inference
            metadata: Model metadata (max_tokens, cost_per_token, etc.)
        """
        self.models[name] = inference_fn
        self.model_metadata[name] = metadata or {}
        self.circuit_breakers[name] = CircuitBreaker()
        self._log(f"Registered LLM model '{name}'")

    def set_active_model(self, name: str):
        """Set the active model."""
        if name in self.models:
            self.active_model = name
            self._log(f"Active LLM model set to '{name}'")
        else:
            self._log(f"Model '{name}' not found")
    
    def set_fallback_chain(self, model_names: List[str]):
        """
        Set fallback chain for model failures.
        
        Args:
            model_names: List of model names in order of preference
        """
        valid_models = [name for name in model_names if name in self.models]
        self.fallback_chain = valid_models
        self._log(f"Set fallback chain: {valid_models}")

    def generate(self, 
                prompt: str, 
                use_cache: bool = True,
                **kwargs) -> Optional[str]:
        """
        Generate response with retry and fallback.
        
        Args:
            prompt: Input prompt
            use_cache: Whether to use response cache
            **kwargs: Additional parameters
            
        Returns:
            Generated response or None
        """
        if not self.active_model:
            self._log("No active model set")
            return None
        
        self.metrics['total_requests'] += 1
        
        # Check cache
        if use_cache and self.enable_caching:
            cache_key = self._get_cache_key(prompt, kwargs)
            if cache_key in self.cache:
                self.metrics['cache_hits'] += 1
                self._log(f"Cache hit for prompt")
                return self.cache[cache_key]
        
        # Try active model with fallback chain
        models_to_try = [self.active_model] + self.fallback_chain
        models_to_try = list(dict.fromkeys(models_to_try))  # Remove duplicates
        
        for model_name in models_to_try:
            result = self._generate_with_retry(model_name, prompt, **kwargs)
            
            if result is not None:
                # Cache successful response
                if self.enable_caching:
                    cache_key = self._get_cache_key(prompt, kwargs)
                    self.cache[cache_key] = result
                
                self.metrics['successful_requests'] += 1
                return result
            
            self._log(f"Model '{model_name}' failed, trying next in chain")
        
        # All models failed
        self.metrics['failed_requests'] += 1
        self._log("All models in chain failed")
        return None
    
    def _generate_with_retry(self, 
                            model_name: str,
                            prompt: str,
                            **kwargs) -> Optional[str]:
        """Generate with exponential backoff retry."""
        circuit_breaker = self.circuit_breakers[model_name]
        stats = self.model_stats[model_name]
        stats['requests'] += 1
        
        for attempt in range(self.max_retries):
            try:
                start_time = time.time()
                
                # Use circuit breaker
                result = circuit_breaker.call(
                    self.models[model_name],
                    prompt,
                    kwargs
                )
                
                # Update metrics
                latency = time.time() - start_time
                stats['successes'] += 1
                stats['total_latency'] += latency
                stats['avg_latency'] = stats['total_latency'] / stats['successes']
                
                # Estimate tokens (rough approximation)
                estimated_tokens = len(prompt.split()) + len(str(result).split())
                self.metrics['total_tokens'] += estimated_tokens
                
                return result
                
            except Exception as e:
                stats['failures'] += 1
                self.metrics['total_retries'] += 1
                
                if attempt < self.max_retries - 1:
                    # Exponential backoff
                    wait_time = 2 ** attempt
                    self._log(f"Attempt {attempt + 1} failed, retrying in {wait_time}s: {e}")
                    time.sleep(wait_time)
                else:
                    self._log(f"All {self.max_retries} attempts failed for model '{model_name}': {e}")
        
        return None
    
    def _get_cache_key(self, prompt: str, kwargs: Dict[str, Any]) -> str:
        """Generate cache key from prompt and parameters."""
        # Create deterministic hash
        cache_string = f"{prompt}:{sorted(kwargs.items())}"
        return hashlib.md5(cache_string.encode()).hexdigest()
    
    def clear_cache(self):
        """Clear response cache."""
        self.cache.clear()
        self._log("Cleared response cache")
    
    def get_model_info(self, model_name: str) -> Optional[Dict[str, Any]]:
        """Get information about a specific model."""
        if model_name not in self.models:
            return None
        
        return {
            'name': model_name,
            'metadata': self.model_metadata.get(model_name, {}),
            'stats': dict(self.model_stats[model_name]),
            'circuit_breaker_state': self.circuit_breakers[model_name].state
        }
    
    def get_metrics(self) -> Dict[str, Any]:
        """Get overall metrics."""
        success_rate = 0.0
        if self.metrics['total_requests'] > 0:
            success_rate = self.metrics['successful_requests'] / self.metrics['total_requests']
        
        cache_hit_rate = 0.0
        if self.metrics['total_requests'] > 0:
            cache_hit_rate = self.metrics['cache_hits'] / self.metrics['total_requests']
        
        return {
            **self.metrics,
            'success_rate': success_rate,
            'cache_hit_rate': cache_hit_rate,
            'active_model': self.active_model,
            'fallback_chain': self.fallback_chain
        }
    
    def reset_circuit_breaker(self, model_name: str):
        """Manually reset circuit breaker for a model."""
        if model_name in self.circuit_breakers:
            cb = self.circuit_breakers[model_name]
            cb.state = "closed"
            cb.failure_count = 0
            cb.last_failure_time = None
            self._log(f"Reset circuit breaker for model '{model_name}'")

    def list_models(self) -> List[str]:
        """List all registered models."""
        return list(self.models.keys())

    def _log(self, message: str):
        """Log a message."""
        print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [LLMManager] {message}")
