"""
Self-Healing Error Recovery System for AIIQ
Provides automatic error classification, retry strategies, and fix suggestions
"""
import re
import asyncio
import logging
from datetime import datetime
from typing import List, Dict, Optional, Any, Callable
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
import json

logger = logging.getLogger(__name__)


class ErrorCategory(Enum):
    """Categories of errors for classification"""
    SYNTAX_ERROR = "syntax"           # Code syntax errors
    RUNTIME_ERROR = "runtime"         # Runtime exceptions
    IMPORT_ERROR = "import"           # Module import failures
    API_ERROR = "api"                 # External API failures
    NETWORK_ERROR = "network"         # Network connectivity issues
    FILE_NOT_FOUND = "file_not_found" # Missing files
    PERMISSION_ERROR = "permission"   # Access denied
    TIMEOUT_ERROR = "timeout"         # Operation timeouts
    LOGIC_ERROR = "logic"             # Logic bugs
    VALIDATION_ERROR = "validation"   # Validation failures
    UNKNOWN_ERROR = "unknown"         # Unclassified


class RecoveryStrategy(Enum):
    """Available recovery strategies"""
    RETRY_SAME = "retry_same"                   # Try again same way
    RETRY_WITH_BACKOFF = "retry_backoff"         # Retry with delay
    RETRY_ALTERNATIVE = "retry_alternative"      # Try different approach
    FALLBACK_ALTERNATIVE = "fallback_alt"       # Use alternative method
    REQUEST_FIX = "request_fix"                 # Ask LLM for fix
    ROLLBACK = "rollback"                       # Undo and retry
    SKIP = "skip"                               # Skip this step
    ABORT = "abort"                             # Give up


@dataclass
class ErrorPattern:
    """A known error pattern with recovery strategy"""
    pattern: str                          # Regex pattern
    category: ErrorCategory
    strategy: RecoveryStrategy
    description: str
    example_fix: str = ""
    max_retries: int = 3


@dataclass
class ErrorRecord:
    """Record of an error and recovery attempt"""
    timestamp: str
    error_message: str
    category: ErrorCategory
    strategy_used: RecoveryStrategy
    attempts: int
    success: bool
    fix_applied: Optional[str] = None
    duration_seconds: float = 0.0


@dataclass
class RecoveryContext:
    """Context for error recovery"""
    original_error: Exception
    category: ErrorCategory
    attempts_made: int = 0
    max_attempts: int = 3
    context: Dict[str, Any] = field(default_factory=dict)
    previous_fixes: List[str] = field(default_factory=list)


class SelfHealingErrorRecovery:
    """
    Self-healing error recovery system.
    
    Features:
    - Error classification (syntax, runtime, logic, API, etc.)
    - Automatic retry with different strategies
    - Error pattern learning (remember what worked)
    - Fix suggestion system
    - Rollback capability
    """
    
    def __init__(self, storage_path: str = "data/error_recovery"):
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(parents=True, exist_ok=True)
        
        # Error patterns database
        self.patterns = self._load_patterns()
        
        # Learning: what strategies worked
        self.strategy_history: Dict[str, List[ErrorRecord]] = {}
        
        # Active recoveries
        self.active_recoveries: Dict[str, RecoveryContext] = {}
        
        logger.info("[Self-Healing] Initialized with %d error patterns", len(self.patterns))
    
    # ========== Error Classification ==========
    
    def classify_error(self, error: Exception) -> ErrorCategory:
        """
        Classify an error into a category.
        
        Args:
            error: The exception to classify
            
        Returns:
            ErrorCategory
        """
        error_str = str(error).lower()
        error_type = type(error).__name__.lower()
        
        # Check against known patterns
        for pattern in self.patterns:
            if re.search(pattern.pattern, error_str, re.IGNORECASE):
                logger.debug("[Self-Healing] Matched pattern: %s", pattern.description)
                return pattern.category
        
        # Type-based classification
        if "syntax" in error_str or "syntaxerror" in error_type:
            return ErrorCategory.SYNTAX_ERROR
        elif "import" in error_str or "module" in error_str:
            return ErrorCategory.IMPORT_ERROR
        elif "timeout" in error_str or "timed out" in error_str:
            return ErrorCategory.TIMEOUT_ERROR
        elif "connection" in error_str or "network" in error_str:
            return ErrorCategory.NETWORK_ERROR
        elif "permission" in error_str or "access denied" in error_str:
            return ErrorCategory.PERMISSION_ERROR
        elif "not found" in error_str or "no such file" in error_str:
            return ErrorCategory.FILE_NOT_FOUND
        elif "api" in error_str or "http" in error_str or "404" in error_str or "500" in error_str:
            return ErrorCategory.API_ERROR
        elif "validation" in error_str or "invalid" in error_str:
            return ErrorCategory.VALIDATION_ERROR
        
        return ErrorCategory.UNKNOWN_ERROR
    
    # ========== Recovery Strategies ==========
    
    async def recover_from_error(
        self,
        error: Exception,
        operation: Callable,
        context: Optional[Dict[str, Any]] = None
    ) -> Any:
        """
        Attempt to recover from an error using various strategies.
        
        Args:
            error: The error that occurred
            operation: The operation to retry
            context: Additional context for recovery
            
        Returns:
            Result of successful recovery or raises original error
        """
        category = self.classify_error(error)
        
        logger.info(
            "[Self-Healing] Error classified as %s: %s",
            category.value, str(error)[:100]
        )
        
        # Get recommended strategy for this category
        strategy = self._get_strategy_for_category(category)
        
        # Create recovery context
        recovery_ctx = RecoveryContext(
            original_error=error,
            category=category,
            max_attempts=3,
            context=context or {}
        )
        
        # Try recovery strategies in order
        for attempt in range(1, 4):  # Max 3 attempts
            recovery_ctx.attempts_made = attempt
            
            try:
                logger.info(
                    "[Self-Healing] Attempt %d/%d using strategy: %s",
                    attempt, 3, strategy.value
                )
                
                result = await self._apply_strategy(
                    strategy,
                    operation,
                    recovery_ctx
                )
                
                # Success! Record what worked
                self._record_success(category, strategy, attempt)
                
                logger.info(
                    "[Self-Healing] ✓ Recovery succeeded on attempt %d",
                    attempt
                )
                
                return result
                
            except Exception as e:
                logger.warning(
                    "[Self-Healing] ✗ Attempt %d failed: %s",
                    attempt, str(e)[:100]
                )
                
                # Try next strategy
                strategy = self._get_next_strategy(strategy)
                
                if strategy == RecoveryStrategy.ABORT:
                    break
        
        # All attempts failed
        self._record_failure(category, recovery_ctx)
        
        # Generate fix suggestion
        suggestion = await self._generate_fix_suggestion(error, category, recovery_ctx)
        
        logger.error("[Self-Healing] All recovery attempts failed")
        logger.error("[Self-Healing] Suggested fix: %s", suggestion)
        
        # Raise with suggestion
        raise type(error)(
            f"{str(error)}\n\n[Suggested Fix]\n{suggestion}"
        ) from error
    
    async def _apply_strategy(
        self,
        strategy: RecoveryStrategy,
        operation: Callable,
        context: RecoveryContext
    ) -> Any:
        """Apply a specific recovery strategy"""
        
        if strategy == RecoveryStrategy.RETRY_SAME:
            # Just try again
            return await operation()
        
        elif strategy == RecoveryStrategy.RETRY_WITH_BACKOFF:
            # Retry with exponential backoff
            delay = min(2 ** context.attempts_made, 10)  # Max 10s
            await asyncio.sleep(delay)
            return await operation()
        
        elif strategy == RecoveryStrategy.RETRY_ALTERNATIVE:
            # Try alternative (modify parameters)
            return await self._try_alternative(operation, context)
        
        elif strategy == RecoveryStrategy.FALLBACK_ALTERNATIVE:
            # Use completely different method
            return await self._use_fallback_method(operation, context)
        
        elif strategy == RecoveryStrategy.REQUEST_FIX:
            # Ask LLM to fix the code
            fixed_code = await self._request_llm_fix(context)
            # Apply the fix and retry
            return await self._apply_fix_and_retry(fixed_code, operation, context)
        
        elif strategy == RecoveryStrategy.SKIP:
            # Skip this operation
            logger.info("[Self-Healing] Skipping operation due to error")
            return None
        
        else:
            # ABORT or unknown
            raise context.original_error
    
    async def _try_alternative(
        self,
        operation: Callable,
        context: RecoveryContext
    ) -> Any:
        """Try with alternative parameters"""
        # Modify context based on error type
        if context.category == ErrorCategory.TIMEOUT_ERROR:
            # Increase timeout
            if "timeout" in context.context:
                context.context["timeout"] *= 2
        
        elif context.category == ErrorCategory.API_ERROR:
            # Try different API endpoint or parameters
            if "retry_count" not in context.context:
                context.context["retry_count"] = 0
            context.context["retry_count"] += 1
        
        return await operation(**context.context)
    
    async def _use_fallback_method(
        self,
        operation: Callable,
        context: RecoveryContext
    ) -> Any:
        """Use a completely different approach"""
        # Check for fallbacks in context
        if "fallback" in context.context:
            fallback_op = context.context["fallback"]
            return await fallback_op()
        
        # No fallback available
        raise context.original_error
    
    async def _request_llm_fix(self, context: RecoveryContext) -> str:
        """Ask LLM to suggest a fix for the error"""
        error_msg = str(context.original_error)
        
        prompt = f"""This error occurred:

{error_msg}

Context: {context.context}

Suggest a fix in 2-3 sentences. Focus on:
1. What's causing the error
2. How to fix it
3. Code changes needed

Fix:"""
        
        try:
            from ai_orchestrator.services.provider_fallback import complete_with_fallback
            
            fix = await complete_with_fallback(
                messages=[{"role": "user", "content": prompt}],
                model="gpt-4o-mini"
            )
            
            if fix:
                if hasattr(fix, 'choices'):
                    return fix.choices[0].message.content
                return str(fix)
        
        except Exception as e:
            logger.warning("[Self-Healing] Could not get LLM fix: %s", e)
        
        return "No automatic fix available. Manual intervention required."
    
    async def _apply_fix_and_retry(
        self,
        fix: str,
        operation: Callable,
        context: RecoveryContext
    ) -> Any:
        """Apply the suggested fix and retry"""
        # Record the fix
        context.previous_fixes.append(fix)
        
        # For now, just log and retry original
        # In a full implementation, would apply code changes
        logger.info("[Self-Healing] Applying fix: %s", fix[:200])
        
        # Retry original operation
        return await operation()
    
    # ========== Strategy Selection ==========
    
    def _get_strategy_for_category(self, category: ErrorCategory) -> RecoveryStrategy:
        """Get recommended recovery strategy for an error category"""
        strategies = {
            ErrorCategory.SYNTAX_ERROR: RecoveryStrategy.REQUEST_FIX,
            ErrorCategory.RUNTIME_ERROR: RecoveryStrategy.RETRY_WITH_BACKOFF,
            ErrorCategory.IMPORT_ERROR: RecoveryStrategy.REQUEST_FIX,
            ErrorCategory.API_ERROR: RecoveryStrategy.RETRY_WITH_BACKOFF,
            ErrorCategory.NETWORK_ERROR: RecoveryStrategy.RETRY_WITH_BACKOFF,
            ErrorCategory.FILE_NOT_FOUND: RecoveryStrategy.REQUEST_FIX,
            ErrorCategory.PERMISSION_ERROR: RecoveryStrategy.ABORT,
            ErrorCategory.TIMEOUT_ERROR: RecoveryStrategy.RETRY_WITH_BACKOFF,
            ErrorCategory.LOGIC_ERROR: RecoveryStrategy.REQUEST_FIX,
            ErrorCategory.VALIDATION_ERROR: RecoveryStrategy.RETRY_ALTERNATIVE,
            ErrorCategory.UNKNOWN_ERROR: RecoveryStrategy.RETRY_SAME,
        }
        
        return strategies.get(category, RecoveryStrategy.RETRY_SAME)
    
    def _get_next_strategy(self, current: RecoveryStrategy) -> RecoveryStrategy:
        """Get next strategy to try if current fails"""
        sequence = {
            RecoveryStrategy.RETRY_SAME: RecoveryStrategy.RETRY_WITH_BACKOFF,
            RecoveryStrategy.RETRY_WITH_BACKOFF: RecoveryStrategy.RETRY_ALTERNATIVE,
            RecoveryStrategy.RETRY_ALTERNATIVE: RecoveryStrategy.FALLBACK_ALTERNATIVE,
            RecoveryStrategy.FALLBACK_ALTERNATIVE: RecoveryStrategy.REQUEST_FIX,
            RecoveryStrategy.REQUEST_FIX: RecoveryStrategy.ABORT,
            RecoveryStrategy.ABORT: RecoveryStrategy.ABORT,
        }
        
        return sequence.get(current, RecoveryStrategy.ABORT)
    
    # ========== Learning & Analytics ==========
    
    def _record_success(
        self,
        category: ErrorCategory,
        strategy: RecoveryStrategy,
        attempt: int
    ) -> None:
        """Record successful recovery for learning"""
        key = f"{category.value}_{strategy.value}"
        
        if key not in self.strategy_history:
            self.strategy_history[key] = []
        
        record = ErrorRecord(
            timestamp=datetime.now().isoformat(),
            error_message="",  # Would be filled in actual use
            category=category,
            strategy_used=strategy,
            attempts=attempt,
            success=True
        )
        
        self.strategy_history[key].append(record)
        
        # Persist to disk
        self._save_history()
    
    def _record_failure(
        self,
        category: ErrorCategory,
        context: RecoveryContext
    ) -> None:
        """Record failed recovery"""
        key = f"{category.value}_failed"
        
        if key not in self.strategy_history:
            self.strategy_history[key] = []
        
        record = ErrorRecord(
            timestamp=datetime.now().isoformat(),
            error_message=str(context.original_error),
            category=category,
            strategy_used=RecoveryStrategy.ABORT,
            attempts=context.attempts_made,
            success=False
        )
        
        self.strategy_history[key].append(record)
        
        self._save_history()
    
    def _load_patterns(self) -> List[ErrorPattern]:
        """Load error patterns from config or use defaults"""
        patterns = [
            # Syntax errors
            ErrorPattern(
                r"SyntaxError|syntax error",
                ErrorCategory.SYNTAX_ERROR,
                RecoveryStrategy.REQUEST_FIX,
                "Code syntax error - needs code fix"
            ),
            
            # Import errors
            ErrorPattern(
                r"ModuleNotFoundError|ImportError|No module named",
                ErrorCategory.IMPORT_ERROR,
                RecoveryStrategy.REQUEST_FIX,
                "Missing import - may need to install package",
                "pip install <package>"
            ),
            
            # API errors
            ErrorPattern(
                r"404|500|502|503|API error|rate limit",
                ErrorCategory.API_ERROR,
                RecoveryStrategy.RETRY_WITH_BACKOFF,
                "External API failure - retry with backoff"
            ),
            
            # Network errors
            ErrorPattern(
                r"Connection refused|Network unreachable|Timeout",
                ErrorCategory.NETWORK_ERROR,
                RecoveryStrategy.RETRY_WITH_BACKOFF,
                "Network issue - retry with backoff"
            ),
            
            # File not found
            ErrorPattern(
                r"FileNotFoundError|No such file|does not exist",
                ErrorCategory.FILE_NOT_FOUND,
                RecoveryStrategy.REQUEST_FIX,
                "File missing - create file or fix path"
            ),
            
            # Permission errors
            ErrorPattern(
                r"Permission denied|Access denied|Unauthorized",
                ErrorCategory.PERMISSION_ERROR,
                RecoveryStrategy.ABORT,
                "Permission issue - cannot auto-recover"
            ),
            
            # Timeout errors
            ErrorPattern(
                r"Timeout|timed out|took too long",
                ErrorCategory.TIMEOUT_ERROR,
                RecoveryStrategy.RETRY_WITH_BACKOFF,
                "Operation timeout - retry with longer timeout"
            ),
            
            # Validation errors
            ErrorPattern(
                r"Validation failed|Invalid|malformed",
                ErrorCategory.VALIDATION_ERROR,
                RecoveryStrategy.RETRY_ALTERNATIVE,
                "Input validation failed - try alternative format"
            ),
        ]
        
        # Try to load from file
        pattern_file = self.storage_path / "error_patterns.json"
        if pattern_file.exists():
            try:
                with open(pattern_file, 'r') as f:
                    data = json.load(f)
                
                # Merge with defaults
                patterns.extend(data)
                
            except Exception as e:
                logger.warning("[Self-Healing] Could not load patterns: %s", e)
        
        return patterns
    
    def _save_history(self) -> None:
        """Save strategy history to disk"""
        history_file = self.storage_path / "strategy_history.json"
        
        try:
            # Convert to serializable format
            serializable = {}
            for key, records in self.strategy_history.items():
                serializable[key] = [
                    {
                        "timestamp": r.timestamp,
                        "error_message": r.error_message,
                        "category": r.category.value,
                        "strategy": r.strategy_used.value,
                        "attempts": r.attempts,
                        "success": r.success
                    }
                    for r in records[-100:]  # Keep last 100
                ]
            
            with open(history_file, 'w') as f:
                json.dump(serializable, f, indent=2)
                
        except Exception as e:
            logger.warning("[Self-Healing] Could not save history: %s", e)
    
    def _load_history(self) -> None:
        """Load strategy history from disk"""
        history_file = self.storage_path / "strategy_history.json"
        
        if not history_file.exists():
            return
        
        try:
            with open(history_file, 'r') as f:
                data = json.load(f)
            
            # Reconstruct ErrorRecords
            for key, records in data.items():
                self.strategy_history[key] = [
                    ErrorRecord(
                        timestamp=r["timestamp"],
                        error_message=r["error_message"],
                        category=ErrorCategory(r["category"]),
                        strategy_used=RecoveryStrategy(r["strategy"]),
                        attempts=r["attempts"],
                        success=r["success"]
                    )
                    for r in records
                ]
            
            logger.info("[Self-Healing] Loaded %d history records", len(data))
            
        except Exception as e:
            logger.warning("[Self-Healing] Could not load history: %s", e)
    
    async def _generate_fix_suggestion(
        self,
        error: Exception,
        category: ErrorCategory,
        context: RecoveryContext
    ) -> str:
        """Generate a fix suggestion for the error"""
        # Get pattern-based suggestion
        for pattern in self.patterns:
            if pattern.category == category and pattern.example_fix:
                if re.search(pattern.pattern, str(error), re.IGNORECASE):
                    return f"{pattern.description}\n\nSuggested: {pattern.example_fix}"
        
        # Generic suggestion
        suggestions = {
            ErrorCategory.SYNTAX_ERROR: "Review code for syntax errors. Check for missing colons, brackets, quotes.",
            ErrorCategory.IMPORT_ERROR: "Install missing package: pip install <package_name>",
            ErrorCategory.API_ERROR: "API service unavailable. Check service status or try again later.",
            ErrorCategory.NETWORK_ERROR: "Network connection issue. Check internet connection.",
            ErrorCategory.FILE_NOT_FOUND: "Create the missing file or verify the path is correct.",
            ErrorCategory.PERMISSION_ERROR: "Check file permissions. Run with appropriate access rights.",
            ErrorCategory.TIMEOUT_ERROR: "Operation timed out. Try increasing timeout or optimizing operation.",
            ErrorCategory.LOGIC_ERROR: "Review code logic. Add debug logging to trace execution.",
            ErrorCategory.VALIDATION_ERROR: "Check input format and required parameters.",
        }
        
        return suggestions.get(
            category,
            "Error occurred. Review logs for details."
        )
    
    # ========== Rollback ==========
    
    async def create_checkpoint(self, name: str) -> str:
        """
        Create a checkpoint for potential rollback.
        
        Args:
            name: Checkpoint name
            
        Returns:
            Checkpoint ID
        """
        checkpoint_id = f"{name}_{datetime.now().isoformat()}"
        
        # In full implementation, would:
        # - Snapshot file system
        # - Save git state
        # - Record variable states
        
        logger.info("[Self-Healing] Created checkpoint: %s", checkpoint_id)
        
        return checkpoint_id
    
    async def rollback_to_checkpoint(self, checkpoint_id: str) -> bool:
        """
        Rollback to a previous checkpoint.
        
        Args:
            checkpoint_id: Checkpoint to rollback to
            
        Returns:
            Success status
        """
        logger.info("[Self-Healing] Rolling back to checkpoint: %s", checkpoint_id)
        
        # In full implementation, would:
        # - Restore file system
        # - Reset git state
        # - Restore variables
        
        return True
    
    # ========== Analytics ==========
    
    def get_recovery_stats(self) -> Dict[str, Any]:
        """Get recovery statistics"""
        total_attempts = sum(len(records) for records in self.strategy_history.values())
        successful = sum(
            len([r for r in records if r.success])
            for records in self.strategy_history.values()
        )
        
        return {
            "total_recovery_attempts": total_attempts,
            "successful_recoveries": successful,
            "success_rate": successful / total_attempts if total_attempts > 0 else 0,
            "categories_seen": len(self.strategy_history),
            "patterns_loaded": len(self.patterns)
        }


# Global instance
_recovery_service: Optional[SelfHealingErrorRecovery] = None


def get_self_healing_service() -> SelfHealingErrorRecovery:
    """Get the global self-healing service"""
    global _recovery_service
    if _recovery_service is None:
        _recovery_service = SelfHealingErrorRecovery()
        _recovery_service._load_history()
    return _recovery_service
