"""
Comprehensive error handling and graceful degradation utilities for PrecipGen.

This module provides robust error handling with clear messages, graceful degradation
strategies, and performance optimization utilities to ensure reliable operation
even when components encounter issues.
"""

import logging
import functools
import traceback
import warnings
from typing import Any, Callable, Dict, List, Optional, Union, Tuple
from datetime import datetime
import pandas as pd
import numpy as np
from pathlib import Path
import json

from .exceptions import (
    PrecipGenError, ValidationError, DataQualityError, ParameterEstimationError,
    SimulationError, StateError, EngineError, handle_graceful_degradation
)

logger = logging.getLogger(__name__)


class ErrorRecoveryManager:
    """
    Manages error recovery and graceful degradation strategies across PrecipGen components.
    
    This class implements comprehensive error handling that allows workflows to continue
    operating even when individual components encounter issues, providing fallback
    strategies and clear guidance for resolution.
    """
    
    def __init__(self):
        self.recovery_strategies = {}
        self.fallback_values = {}
        self.error_history = []
        self.degradation_warnings = []
    
    def register_recovery_strategy(self, error_type: type, strategy: Callable) -> None:
        """Register a recovery strategy for a specific error type."""
        self.recovery_strategies[error_type] = strategy
        logger.debug(f"Registered recovery strategy for {error_type.__name__}")
    
    def register_fallback_value(self, operation: str, fallback: Any) -> None:
        """Register a fallback value for a specific operation."""
        self.fallback_values[operation] = fallback
        logger.debug(f"Registered fallback value for operation: {operation}")
    
    def handle_error_with_recovery(self, error: Exception, operation: str, 
                                 context: Optional[Dict[str, Any]] = None) -> Tuple[bool, Any]:
        """
        Handle error with appropriate recovery strategy.
        
        Args:
            error: The exception that occurred
            operation: Name of the operation that failed
            context: Additional context for recovery
            
        Returns:
            Tuple of (recovery_successful, result_or_fallback)
        """
        self.error_history.append({
            'timestamp': datetime.now().isoformat(),
            'error_type': type(error).__name__,
            'operation': operation,
            'message': str(error),
            'context': context or {}
        })
        
        # Try registered recovery strategy first
        error_type = type(error)
        if error_type in self.recovery_strategies:
            try:
                result = self.recovery_strategies[error_type](error, operation, context)
                logger.info(f"Successfully recovered from {error_type.__name__} in {operation}")
                return True, result
            except Exception as recovery_error:
                logger.warning(f"Recovery strategy failed: {recovery_error}")
        
        # Try fallback value
        if operation in self.fallback_values:
            fallback = self.fallback_values[operation]
            warning_msg = f"Using fallback value for {operation} due to {error_type.__name__}: {error}"
            self.degradation_warnings.append(warning_msg)
            warnings.warn(warning_msg, UserWarning)
            logger.warning(warning_msg)
            return True, fallback
        
        # No recovery possible
        logger.error(f"No recovery strategy available for {error_type.__name__} in {operation}")
        return False, None
    
    def get_error_summary(self) -> Dict[str, Any]:
        """Get summary of errors and recovery actions."""
        return {
            'total_errors': len(self.error_history),
            'error_types': list(set(e['error_type'] for e in self.error_history)),
            'degradation_warnings': len(self.degradation_warnings),
            'recent_errors': self.error_history[-5:] if self.error_history else [],
            'recovery_strategies_available': len(self.recovery_strategies),
            'fallback_values_available': len(self.fallback_values)
        }


def robust_operation(operation_name: str, 
                    fallback_value: Any = None,
                    error_types: Tuple[type, ...] = (Exception,),
                    log_errors: bool = True,
                    raise_on_failure: bool = False):
    """
    Decorator for robust operation execution with error handling and graceful degradation.
    
    Args:
        operation_name: Name of the operation for logging
        fallback_value: Value to return if operation fails
        error_types: Tuple of exception types to catch
        log_errors: Whether to log errors
        raise_on_failure: Whether to raise exception after logging
        
    Returns:
        Decorated function with error handling
    """
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except error_types as e:
                if log_errors:
                    logger.error(f"Error in {operation_name}: {str(e)}")
                    if logger.isEnabledFor(logging.DEBUG):
                        logger.debug(f"Traceback: {traceback.format_exc()}")
                
                if raise_on_failure:
                    raise
                
                if fallback_value is not None:
                    warning_msg = f"Using fallback value for {operation_name}: {fallback_value}"
                    warnings.warn(warning_msg, UserWarning)
                    logger.warning(warning_msg)
                    return fallback_value
                
                # Return None if no fallback specified
                return None
        
        return wrapper
    return decorator


def validate_with_recovery(data: Any, 
                         validation_func: Callable,
                         recovery_func: Optional[Callable] = None,
                         operation_name: str = "validation") -> Tuple[bool, Any, List[str]]:
    """
    Validate data with optional recovery function.
    
    Args:
        data: Data to validate
        validation_func: Function that validates data and returns (is_valid, errors)
        recovery_func: Optional function to attempt data recovery
        operation_name: Name for logging
        
    Returns:
        Tuple of (is_valid, processed_data, error_messages)
    """
    try:
        is_valid, errors = validation_func(data)
        
        if is_valid:
            return True, data, []
        
        # Validation failed - try recovery if available
        if recovery_func is not None:
            try:
                recovered_data = recovery_func(data, errors)
                # Re-validate recovered data
                is_recovered_valid, recovery_errors = validation_func(recovered_data)
                
                if is_recovered_valid:
                    logger.info(f"Successfully recovered data in {operation_name}")
                    return True, recovered_data, [f"Data recovered from: {', '.join(errors)}"]
                else:
                    logger.warning(f"Data recovery failed in {operation_name}")
                    return False, data, errors + recovery_errors
            
            except Exception as recovery_error:
                logger.error(f"Recovery function failed in {operation_name}: {recovery_error}")
                return False, data, errors + [f"Recovery failed: {str(recovery_error)}"]
        
        return False, data, errors
    
    except Exception as e:
        error_msg = f"Validation process failed in {operation_name}: {str(e)}"
        logger.error(error_msg)
        return False, data, [error_msg]


class PerformanceOptimizer:
    """
    Performance optimization utilities for PrecipGen workflows.
    
    Provides caching, memory management, and performance monitoring
    to optimize workflow execution and reduce computational overhead.
    """
    
    def __init__(self):
        self.cache = {}
        self.cache_metadata = {}
        self.performance_metrics = {}
        self.memory_usage = {}
    
    def cached_computation(self, cache_key: str, 
                         computation_func: Callable,
                         *args, **kwargs) -> Any:
        """
        Execute computation with caching to avoid redundant calculations.
        
        Args:
            cache_key: Unique key for caching
            computation_func: Function to execute
            *args, **kwargs: Arguments for the function
            
        Returns:
            Computation result (cached or newly computed)
        """
        if cache_key in self.cache:
            logger.debug(f"Cache hit for key: {cache_key}")
            self.cache_metadata[cache_key]['access_count'] += 1
            self.cache_metadata[cache_key]['last_accessed'] = datetime.now().isoformat()
            return self.cache[cache_key]
        
        # Compute and cache result
        start_time = datetime.now()
        try:
            result = computation_func(*args, **kwargs)
            computation_time = (datetime.now() - start_time).total_seconds()
            
            # Store in cache with metadata
            self.cache[cache_key] = result
            self.cache_metadata[cache_key] = {
                'created_at': start_time.isoformat(),
                'computation_time': computation_time,
                'access_count': 1,
                'last_accessed': start_time.isoformat(),
                'result_type': type(result).__name__
            }
            
            logger.debug(f"Cached computation result for key: {cache_key} (time: {computation_time:.3f}s)")
            return result
            
        except Exception as e:
            logger.error(f"Cached computation failed for key {cache_key}: {e}")
            raise
    
    def optimize_data_operations(self, data: pd.Series, 
                               operations: List[str]) -> Dict[str, Any]:
        """
        Optimize data operations by batching and caching intermediate results.
        
        Args:
            data: Input data series
            operations: List of operations to perform
            
        Returns:
            Dictionary of operation results
        """
        results = {}
        
        # Create data signature for caching
        data_signature = self._create_data_signature(data)
        
        for operation in operations:
            cache_key = f"{data_signature}_{operation}"
            
            if operation == 'basic_stats':
                results[operation] = self.cached_computation(
                    cache_key, self._compute_basic_stats, data
                )
            elif operation == 'monthly_aggregation':
                results[operation] = self.cached_computation(
                    cache_key, self._compute_monthly_aggregation, data
                )
            elif operation == 'annual_aggregation':
                results[operation] = self.cached_computation(
                    cache_key, self._compute_annual_aggregation, data
                )
            elif operation == 'wet_day_analysis':
                results[operation] = self.cached_computation(
                    cache_key, self._compute_wet_day_analysis, data
                )
            else:
                logger.warning(f"Unknown operation: {operation}")
        
        return results
    
    def _create_data_signature(self, data: pd.Series) -> str:
        """Create a signature for data caching."""
        return f"{len(data)}_{hash(tuple(data.values[:min(100, len(data))]))}"
    
    def _compute_basic_stats(self, data: pd.Series) -> Dict[str, float]:
        """Compute basic statistics."""
        return {
            'mean': float(data.mean()),
            'std': float(data.std()),
            'min': float(data.min()),
            'max': float(data.max()),
            'sum': float(data.sum()),
            'count': len(data)
        }
    
    def _compute_monthly_aggregation(self, data: pd.Series) -> Dict[int, Dict[str, float]]:
        """Compute monthly aggregations."""
        if not isinstance(data.index, pd.DatetimeIndex):
            return {}
        
        monthly_stats = {}
        for month in range(1, 13):
            month_data = data[data.index.month == month]
            if len(month_data) > 0:
                monthly_stats[month] = self._compute_basic_stats(month_data)
        
        return monthly_stats
    
    def _compute_annual_aggregation(self, data: pd.Series) -> Dict[int, Dict[str, float]]:
        """Compute annual aggregations."""
        if not isinstance(data.index, pd.DatetimeIndex):
            return {}
        
        annual_stats = {}
        for year in data.index.year.unique():
            year_data = data[data.index.year == year]
            if len(year_data) > 0:
                annual_stats[year] = self._compute_basic_stats(year_data)
        
        return annual_stats
    
    def _compute_wet_day_analysis(self, data: pd.Series, 
                                threshold: float = 0.0254) -> Dict[str, Any]:
        """Compute wet day analysis."""
        wet_days = data > threshold
        return {
            'wet_day_count': int(wet_days.sum()),
            'wet_day_percentage': float(wet_days.mean() * 100),
            'mean_wet_day_amount': float(data[wet_days].mean()) if wet_days.any() else 0.0,
            'max_wet_day_amount': float(data[wet_days].max()) if wet_days.any() else 0.0
        }
    
    def clear_cache(self, pattern: Optional[str] = None) -> None:
        """Clear cache entries, optionally matching a pattern."""
        if pattern is None:
            cleared_count = len(self.cache)
            self.cache.clear()
            self.cache_metadata.clear()
            logger.info(f"Cleared all cache entries ({cleared_count} items)")
        else:
            cleared_keys = [key for key in self.cache.keys() if pattern in key]
            for key in cleared_keys:
                del self.cache[key]
                del self.cache_metadata[key]
            logger.info(f"Cleared {len(cleared_keys)} cache entries matching pattern: {pattern}")
    
    def get_cache_info(self) -> Dict[str, Any]:
        """Get cache performance information."""
        total_access_count = sum(
            meta['access_count'] for meta in self.cache_metadata.values()
        )
        
        return {
            'total_entries': len(self.cache),
            'total_accesses': total_access_count,
            'cache_hit_ratio': total_access_count / max(1, len(self.cache)),
            'memory_estimate_mb': sum(
                len(str(value)) for value in self.cache.values()
            ) / (1024 * 1024),
            'oldest_entry': min(
                (meta['created_at'] for meta in self.cache_metadata.values()),
                default=None
            ),
            'most_accessed_key': max(
                self.cache_metadata.items(),
                key=lambda x: x[1]['access_count'],
                default=(None, {'access_count': 0})
            )[0]
        }


class GracefulDegradationManager:
    """
    Manages graceful degradation strategies for PrecipGen components.
    
    When components encounter errors, this manager provides fallback strategies
    that allow workflows to continue with reduced functionality rather than
    complete failure.
    """
    
    def __init__(self):
        self.degradation_strategies = {}
        self.active_degradations = {}
        self.degradation_history = []
    
    def register_degradation_strategy(self, component: str, 
                                    strategy: Callable) -> None:
        """Register a degradation strategy for a component."""
        self.degradation_strategies[component] = strategy
        logger.debug(f"Registered degradation strategy for component: {component}")
    
    def apply_degradation(self, component: str, error: Exception, 
                        context: Dict[str, Any]) -> Any:
        """
        Apply graceful degradation for a component.
        
        Args:
            component: Name of the component
            error: The error that triggered degradation
            context: Context information for degradation
            
        Returns:
            Degraded result or fallback value
        """
        if component not in self.degradation_strategies:
            logger.error(f"No degradation strategy available for component: {component}")
            raise error
        
        try:
            strategy = self.degradation_strategies[component]
            result = strategy(error, context)
            
            # Record degradation
            degradation_record = {
                'timestamp': datetime.now().isoformat(),
                'component': component,
                'error_type': type(error).__name__,
                'error_message': str(error),
                'context': context,
                'degradation_applied': True
            }
            
            self.degradation_history.append(degradation_record)
            self.active_degradations[component] = degradation_record
            
            warning_msg = f"Applied graceful degradation for {component}: {type(error).__name__}"
            warnings.warn(warning_msg, UserWarning)
            logger.warning(warning_msg)
            
            return result
            
        except Exception as degradation_error:
            logger.error(f"Degradation strategy failed for {component}: {degradation_error}")
            raise error
    
    def is_degraded(self, component: str) -> bool:
        """Check if a component is currently in degraded mode."""
        return component in self.active_degradations
    
    def clear_degradation(self, component: str) -> None:
        """Clear degradation status for a component."""
        if component in self.active_degradations:
            del self.active_degradations[component]
            logger.info(f"Cleared degradation status for component: {component}")
    
    def get_degradation_status(self) -> Dict[str, Any]:
        """Get current degradation status."""
        return {
            'active_degradations': len(self.active_degradations),
            'degraded_components': list(self.active_degradations.keys()),
            'total_degradation_events': len(self.degradation_history),
            'recent_degradations': self.degradation_history[-5:] if self.degradation_history else []
        }


# Global instances for use across PrecipGen
error_recovery_manager = ErrorRecoveryManager()
performance_optimizer = PerformanceOptimizer()
degradation_manager = GracefulDegradationManager()


# Register default recovery strategies
def _default_parameter_recovery(error: Exception, operation: str, context: Dict[str, Any]) -> Any:
    """Default recovery strategy for parameter estimation errors."""
    if 'fallback_parameters' in context:
        return context['fallback_parameters']
    
    # Use generic monthly parameters as last resort
    from ..engines.analytical import MonthlyParams
    return {month: MonthlyParams(p_ww=0.5, p_wd=0.3, alpha=1.0, beta=5.0) 
            for month in range(1, 13)}


def _default_data_recovery(error: Exception, operation: str, context: Dict[str, Any]) -> Any:
    """Default recovery strategy for data loading errors."""
    if 'fallback_data' in context:
        return context['fallback_data']
    
    # Return empty series with proper structure
    return pd.Series([], dtype=float, name='precipitation')


# Register default strategies
error_recovery_manager.register_recovery_strategy(ParameterEstimationError, _default_parameter_recovery)
error_recovery_manager.register_recovery_strategy(ValidationError, _default_data_recovery)

# Register default fallback values
error_recovery_manager.register_fallback_value('parameter_estimation', 'default_parameters')
error_recovery_manager.register_fallback_value('data_loading', 'empty_series')
error_recovery_manager.register_fallback_value('trend_analysis', None)