"""
Intelligent Error Recovery System for Aperion Pipelines.

Provides automatic recovery strategies for pipeline stage failures:
- Retry with exponential backoff
- LLM-powered error diagnosis and remediation hints
- Fallback strategies
- Circuit breaker pattern
"""

import logging
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Optional

from .pipeline import PipelineContext, PipelineStage

logger = logging.getLogger(__name__)


class RecoveryStrategy(Enum):
    """Error recovery strategies."""

    NONE = "none"  # No recovery, fail immediately
    RETRY_IMMEDIATE = "retry_immediate"  # Retry immediately
    RETRY_BACKOFF = "retry_backoff"  # Retry with exponential backoff
    RETRY_WITH_LLM_HINTS = "retry_with_llm_hints"  # Use LLM to diagnose and retry
    FALLBACK = "fallback"  # Use fallback handler
    CIRCUIT_BREAKER = "circuit_breaker"  # Fail fast after threshold


class FailureType(Enum):
    """Types of failures that can occur."""

    TIMEOUT = "timeout"  # Stage execution timeout
    EXCEPTION = "exception"  # General exception
    VALIDATION_ERROR = "validation_error"  # Input/output validation failed
    RESOURCE_ERROR = "resource_error"  # Resource unavailable (network, DB, etc.)
    LLM_ERROR = "llm_error"  # LLM provider failure
    UNKNOWN = "unknown"  # Unknown failure type


@dataclass
class RecoveryConfig:
    """Configuration for error recovery."""

    strategy: RecoveryStrategy = RecoveryStrategy.RETRY_BACKOFF
    max_retries: int = 3
    initial_backoff: float = 1.0  # seconds
    backoff_multiplier: float = 2.0
    max_backoff: float = 60.0  # seconds
    use_llm_diagnosis: bool = True
    fallback_handler: Optional[Callable[[PipelineContext], PipelineContext]] = None
    circuit_breaker_threshold: int = 5  # failures before circuit opens
    circuit_breaker_timeout: float = 60.0  # seconds before circuit reset
    retry_on_failure_types: list[FailureType] = field(
        default_factory=lambda: [
            FailureType.TIMEOUT,
            FailureType.RESOURCE_ERROR,
            FailureType.LLM_ERROR,
            FailureType.EXCEPTION,  # General exceptions are retryable
        ]
    )


@dataclass
class RecoveryAttempt:
    """Record of a recovery attempt."""

    attempt_number: int
    strategy: RecoveryStrategy
    failure_type: FailureType
    error_message: str
    timestamp: float
    backoff_delay: float = 0.0
    llm_hints: Optional[dict[str, Any]] = None
    success: bool = False
    recovery_time: float = 0.0


@dataclass
class RecoveryResult:
    """Result of recovery attempts."""

    success: bool
    attempts: list[RecoveryAttempt]
    final_error: Optional[str] = None
    total_recovery_time: float = 0.0
    strategy_used: Optional[RecoveryStrategy] = None
    context: Optional[PipelineContext] = None


class ErrorClassifier:
    """Classifies errors to determine appropriate recovery strategy."""

    @staticmethod
    def classify(error: Exception) -> FailureType:
        """
        Classify an error to determine its type.

        Args:
            error: Exception that occurred

        Returns:
            Classified failure type
        """
        error_str = str(error).lower()
        error_type = type(error).__name__.lower()

        # Timeout errors (highest priority)
        if "timeout" in error_str or isinstance(error, TimeoutError):
            return FailureType.TIMEOUT

        # LLM errors (check before resource errors to catch "LLM provider unavailable")
        llm_keywords = [
            "llm",
            "provider",
            "model",
            "api key",
            "quota",
            "rate limit",
            "api",
        ]
        if any(keyword in error_str for keyword in llm_keywords):
            return FailureType.LLM_ERROR

        # Resource errors (network, database, file system)
        resource_keywords = [
            "connection",
            "network",
            "unavailable",
            "unreachable",
            "refused",
            "database",
            "file not found",
            "permission denied",
        ]
        if any(keyword in error_str for keyword in resource_keywords):
            return FailureType.RESOURCE_ERROR

        # Validation errors
        validation_keywords = ["validation", "invalid", "malformed", "schema"]
        if any(keyword in error_str for keyword in validation_keywords):
            return FailureType.VALIDATION_ERROR

        # Default to exception (which is retryable by default)
        if "exception" in error_type:
            return FailureType.EXCEPTION

        return FailureType.EXCEPTION  # Default to retryable exception

    @staticmethod
    def is_retryable(failure_type: FailureType, config: RecoveryConfig) -> bool:
        """
        Check if a failure type is retryable based on configuration.

        Args:
            failure_type: Type of failure
            config: Recovery configuration

        Returns:
            True if failure should be retried
        """
        return failure_type in config.retry_on_failure_types


class RecoveryEngine:
    """Engine for executing recovery strategies."""

    def __init__(self):
        """Initialize recovery engine."""
        self.circuit_breaker_state: dict[str, dict[str, Any]] = {}
        self.classifier = ErrorClassifier()

    def recover_stage(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        error: Exception,
        config: RecoveryConfig,
    ) -> RecoveryResult:
        """
        Attempt to recover from a stage failure.

        Args:
            stage: Failed stage
            context: Pipeline context
            error: Exception that occurred
            config: Recovery configuration

        Returns:
            Recovery result with status and details
        """
        start_time = time.time()
        attempts: list[RecoveryAttempt] = []

        # Classify the error
        failure_type = self.classifier.classify(error)
        logger.info(f"Stage {stage.id} failed with {failure_type.value} error: {error}")

        # Check circuit breaker
        if config.strategy == RecoveryStrategy.CIRCUIT_BREAKER:
            if self._is_circuit_open(stage.id, config):
                logger.warning(
                    f"Circuit breaker open for stage {stage.id}, failing fast"
                )
                return RecoveryResult(
                    success=False,
                    attempts=attempts,
                    final_error="Circuit breaker open",
                    total_recovery_time=time.time() - start_time,
                    strategy_used=RecoveryStrategy.CIRCUIT_BREAKER,
                )

        # Check if error is retryable
        if not self.classifier.is_retryable(failure_type, config):
            logger.info(
                f"Failure type {failure_type.value} is not retryable, failing immediately"
            )
            return RecoveryResult(
                success=False,
                attempts=attempts,
                final_error=str(error),
                total_recovery_time=time.time() - start_time,
                strategy_used=RecoveryStrategy.NONE,
            )

        # Attempt recovery based on strategy
        if config.strategy == RecoveryStrategy.RETRY_IMMEDIATE:
            return self._retry_immediate(
                stage, context, error, failure_type, config, attempts
            )

        elif config.strategy == RecoveryStrategy.RETRY_BACKOFF:
            return self._retry_with_backoff(
                stage, context, error, failure_type, config, attempts, start_time
            )

        elif config.strategy == RecoveryStrategy.RETRY_WITH_LLM_HINTS:
            return self._retry_with_llm_hints(
                stage, context, error, failure_type, config, attempts, start_time
            )

        elif config.strategy == RecoveryStrategy.FALLBACK:
            return self._use_fallback(
                stage, context, error, failure_type, config, attempts, start_time
            )

        # Default: no recovery
        return RecoveryResult(
            success=False,
            attempts=attempts,
            final_error=str(error),
            total_recovery_time=time.time() - start_time,
            strategy_used=RecoveryStrategy.NONE,
        )

    def _retry_immediate(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        error: Exception,
        failure_type: FailureType,
        config: RecoveryConfig,
        attempts: list[RecoveryAttempt],
    ) -> RecoveryResult:
        """Retry immediately without delay."""
        start_time = time.time()

        for attempt_num in range(1, config.max_retries + 1):
            attempt = RecoveryAttempt(
                attempt_number=attempt_num,
                strategy=RecoveryStrategy.RETRY_IMMEDIATE,
                failure_type=failure_type,
                error_message=str(error),
                timestamp=time.time(),
            )

            try:
                logger.info(
                    f"Retry attempt {attempt_num}/{config.max_retries} for stage {stage.id}"
                )
                result_context = stage.handler(context)

                # Success!
                attempt.success = True
                attempt.recovery_time = time.time() - attempt.timestamp
                attempts.append(attempt)

                logger.info(f"Stage {stage.id} recovered after {attempt_num} attempts")
                return RecoveryResult(
                    success=True,
                    attempts=attempts,
                    total_recovery_time=time.time() - start_time,
                    strategy_used=RecoveryStrategy.RETRY_IMMEDIATE,
                    context=result_context,
                )

            except Exception as e:
                attempt.recovery_time = time.time() - attempt.timestamp
                attempts.append(attempt)
                error = e
                logger.warning(
                    f"Retry attempt {attempt_num} failed for stage {stage.id}: {e}"
                )

        return RecoveryResult(
            success=False,
            attempts=attempts,
            final_error=str(error),
            total_recovery_time=time.time() - start_time,
            strategy_used=RecoveryStrategy.RETRY_IMMEDIATE,
        )

    def _retry_with_backoff(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        error: Exception,
        failure_type: FailureType,
        config: RecoveryConfig,
        attempts: list[RecoveryAttempt],
        start_time: float,
    ) -> RecoveryResult:
        """Retry with exponential backoff."""
        backoff = config.initial_backoff

        for attempt_num in range(1, config.max_retries + 1):
            # Calculate backoff delay
            delay = min(backoff, config.max_backoff)

            attempt = RecoveryAttempt(
                attempt_number=attempt_num,
                strategy=RecoveryStrategy.RETRY_BACKOFF,
                failure_type=failure_type,
                error_message=str(error),
                timestamp=time.time(),
                backoff_delay=delay,
            )

            # Wait before retry
            if delay > 0:
                logger.info(
                    f"Waiting {delay:.1f}s before retry {attempt_num}/{config.max_retries}"
                )
                time.sleep(delay)

            try:
                logger.info(
                    f"Retry attempt {attempt_num}/{config.max_retries} for stage {stage.id}"
                )
                result_context = stage.handler(context)

                # Success!
                attempt.success = True
                attempt.recovery_time = time.time() - attempt.timestamp
                attempts.append(attempt)

                logger.info(
                    f"Stage {stage.id} recovered after {attempt_num} attempts with backoff"
                )
                return RecoveryResult(
                    success=True,
                    attempts=attempts,
                    total_recovery_time=time.time() - start_time,
                    strategy_used=RecoveryStrategy.RETRY_BACKOFF,
                    context=result_context,
                )

            except Exception as e:
                attempt.recovery_time = time.time() - attempt.timestamp
                attempts.append(attempt)
                error = e
                logger.warning(
                    f"Retry attempt {attempt_num} with backoff failed for stage {stage.id}: {e}"
                )

            # Increase backoff for next attempt
            backoff *= config.backoff_multiplier

        return RecoveryResult(
            success=False,
            attempts=attempts,
            final_error=str(error),
            total_recovery_time=time.time() - start_time,
            strategy_used=RecoveryStrategy.RETRY_BACKOFF,
        )

    def _retry_with_llm_hints(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        error: Exception,
        failure_type: FailureType,
        config: RecoveryConfig,
        attempts: list[RecoveryAttempt],
        start_time: float,
    ) -> RecoveryResult:
        """
        Retry with LLM-powered error diagnosis and hints.

        This strategy uses LLM to analyze the error and provide
        remediation hints before retrying.
        """
        logger.info(f"LLM-assisted recovery for stage {stage.id}")

        # Get LLM diagnosis if enabled
        llm_hints = None
        if config.use_llm_diagnosis:
            llm_hints = self._get_llm_diagnosis(
                stage=stage,
                context=context,
                error=error,
                failure_type=failure_type,
            )

        # Store LLM hints in context for retry attempts
        if llm_hints:
            context.metadata["recovery_hints"] = llm_hints
            logger.info(
                f"LLM diagnosis for {stage.id}: {llm_hints.get('diagnosis', 'No diagnosis')[:100]}"
            )
        else:
            # Fallback hints if LLM unavailable
            context.metadata["recovery_hints"] = {
                "error_type": failure_type.value,
                "error_message": str(error),
                "suggested_action": "Retry with exponential backoff",
                "llm_available": False,
            }

        # Retry with exponential backoff, recording LLM attempts
        retry_count = 0
        backoff_delay = config.initial_backoff
        classifier = ErrorClassifier()

        while retry_count < config.max_retries:
            time.sleep(backoff_delay)

            # Record attempt with LLM strategy
            attempt = RecoveryAttempt(
                attempt_number=len(attempts) + 1,
                strategy=RecoveryStrategy.RETRY_WITH_LLM_HINTS,  # Mark as LLM-assisted
                failure_type=failure_type,
                error_message=str(error),
                timestamp=time.time(),
                backoff_delay=backoff_delay,
                llm_hints=llm_hints,  # Include LLM hints in attempt
                success=False,
            )
            attempts.append(attempt)

            # Retry execution
            try:
                retry_start = time.time()
                result_context = stage.handler(context)
                attempt.recovery_time = time.time() - retry_start
                attempt.success = True

                logger.info(
                    f"LLM-assisted recovery succeeded for {stage.id} on attempt {retry_count + 1}"
                )
                return RecoveryResult(
                    success=True,
                    attempts=attempts,
                    strategy_used=RecoveryStrategy.RETRY_WITH_LLM_HINTS,
                    total_recovery_time=time.time() - start_time,
                    context=result_context,
                )

            except Exception as retry_error:
                attempt.recovery_time = time.time() - retry_start
                logger.warning(
                    f"LLM-assisted retry {retry_count + 1} failed for {stage.id}: {retry_error}"
                )

                # Update error for next attempt
                error = retry_error
                failure_type = classifier.classify(retry_error)

            retry_count += 1
            backoff_delay = min(
                backoff_delay * config.backoff_multiplier, config.max_backoff
            )

        # All retries exhausted
        return RecoveryResult(
            success=False,
            attempts=attempts,
            final_error=str(error),
            strategy_used=RecoveryStrategy.RETRY_WITH_LLM_HINTS,
            total_recovery_time=time.time() - start_time,
            context=None,
        )

    def _get_llm_diagnosis(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        error: Exception,
        failure_type: FailureType,
    ) -> Optional[dict[str, Any]]:
        """
        Use LLM to diagnose error and suggest recovery actions.

        Returns None if LLM providers are unavailable (standalone mode).
        """
        try:
            # Try to import monolith LLM providers (only available when
            # running inside the full Aperion stack)
            from stack.aperion.foundation.llm import load_provider, safe_provider_call
        except ImportError:
            logger.debug("LLM providers not available in standalone mode")
            return None

        try:
            provider = load_provider("openai", fallback=True)
        except Exception as load_error:
            logger.warning(f"Failed to load LLM provider: {load_error}")
            return None

        if provider is None:
            logger.warning("LLM provider unavailable for error diagnosis")
            return None

        try:
            prompt = self._build_diagnosis_prompt(stage, context, error, failure_type)
            result, call_warnings = safe_provider_call(
                provider,
                "chat",
                prompt,
                temperature=0.3,
                max_tokens=500,
            )

            if result is None or "replies" not in result:
                logger.warning("LLM diagnosis failed")
                return None

            diagnosis_text = result["replies"][0] if result["replies"] else ""

            return {
                "diagnosis": diagnosis_text,
                "error_type": failure_type.value,
                "error_message": str(error),
                "suggested_action": self._extract_action_from_diagnosis(diagnosis_text),
                "llm_provider": provider.name,
                "warnings": call_warnings,
            }

        except Exception as e:
            logger.warning(f"Failed to get LLM diagnosis: {e}")
            return None

    def _build_diagnosis_prompt(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        error: Exception,
        failure_type: FailureType,
    ) -> str:
        """Build prompt for LLM error diagnosis."""
        return f"""You are an expert system reliability engineer analyzing a pipeline stage failure.

Stage Information:
- Stage ID: {stage.id}
- Stage Name: {stage.name}
- Failure Type: {failure_type.value}

Error Details:
- Exception Type: {type(error).__name__}
- Error Message: {str(error)}

Context Data Keys: {list(context.data.keys()) if context.data else 'None'}
Context Metadata Keys: {list(context.metadata.keys()) if context.metadata else 'None'}

Task: Analyze this error and provide:
1. Root cause analysis (1-2 sentences)
2. Specific recovery action (be concrete and actionable)
3. Whether retrying is likely to succeed (yes/no/maybe)

Format your response as:
DIAGNOSIS: <root cause>
ACTION: <specific recovery action>
RETRY_RECOMMENDED: <yes/no/maybe>
REASONING: <why this approach should work>

Keep response concise and technical."""

    def _extract_action_from_diagnosis(self, diagnosis_text: str) -> str:
        """Extract actionable recovery suggestion from LLM diagnosis."""
        # Look for ACTION: line in diagnosis
        for line in diagnosis_text.split("\n"):
            if line.strip().startswith("ACTION:"):
                return line.split("ACTION:", 1)[1].strip()

        # Fallback: return first non-empty line
        for line in diagnosis_text.split("\n"):
            if line.strip():
                return line.strip()[:200]  # Truncate long lines

        return "Retry with exponential backoff"

    def _use_fallback(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        error: Exception,
        failure_type: FailureType,
        config: RecoveryConfig,
        attempts: list[RecoveryAttempt],
        start_time: float,
    ) -> RecoveryResult:
        """Use fallback handler instead of retrying."""
        if not config.fallback_handler:
            logger.warning(
                "Fallback strategy requested but no fallback handler configured"
            )
            return RecoveryResult(
                success=False,
                attempts=attempts,
                final_error="No fallback handler configured",
                total_recovery_time=time.time() - start_time,
                strategy_used=RecoveryStrategy.FALLBACK,
            )

        attempt = RecoveryAttempt(
            attempt_number=1,
            strategy=RecoveryStrategy.FALLBACK,
            failure_type=failure_type,
            error_message=str(error),
            timestamp=time.time(),
        )

        try:
            logger.info(f"Using fallback handler for stage {stage.id}")
            result_context = config.fallback_handler(context)

            attempt.success = True
            attempt.recovery_time = time.time() - attempt.timestamp
            attempts.append(attempt)

            logger.info(f"Stage {stage.id} recovered using fallback handler")
            return RecoveryResult(
                success=True,
                attempts=attempts,
                total_recovery_time=time.time() - start_time,
                strategy_used=RecoveryStrategy.FALLBACK,
                context=result_context,
            )

        except Exception as e:
            attempt.recovery_time = time.time() - attempt.timestamp
            attempts.append(attempt)

            logger.error(f"Fallback handler failed for stage {stage.id}: {e}")
            return RecoveryResult(
                success=False,
                attempts=attempts,
                final_error=f"Fallback handler failed: {e}",
                total_recovery_time=time.time() - start_time,
                strategy_used=RecoveryStrategy.FALLBACK,
            )

    def _is_circuit_open(self, stage_id: str, config: RecoveryConfig) -> bool:
        """Check if circuit breaker is open for a stage."""
        if stage_id not in self.circuit_breaker_state:
            self.circuit_breaker_state[stage_id] = {
                "failure_count": 0,
                "last_failure_time": 0,
                "state": "closed",
            }

        state = self.circuit_breaker_state[stage_id]

        # Check if circuit should reset
        if state["state"] == "open":
            time_since_open = time.time() - state["last_failure_time"]
            if time_since_open >= config.circuit_breaker_timeout:
                logger.info(
                    f"Circuit breaker reset for stage {stage_id} after {time_since_open:.1f}s"
                )
                state["state"] = "half-open"
                state["failure_count"] = 0
                return False

            return True

        return False

    def record_failure(self, stage_id: str, config: RecoveryConfig) -> None:
        """Record a failure for circuit breaker tracking."""
        if stage_id not in self.circuit_breaker_state:
            self.circuit_breaker_state[stage_id] = {
                "failure_count": 0,
                "last_failure_time": 0,
                "state": "closed",
            }

        state = self.circuit_breaker_state[stage_id]
        state["failure_count"] += 1
        state["last_failure_time"] = time.time()

        # Open circuit if threshold exceeded
        if state["failure_count"] >= config.circuit_breaker_threshold:
            logger.warning(
                f"Circuit breaker opened for stage {stage_id} after "
                f"{state['failure_count']} failures"
            )
            state["state"] = "open"

    def record_success(self, stage_id: str) -> None:
        """Record a success for circuit breaker tracking."""
        if stage_id in self.circuit_breaker_state:
            state = self.circuit_breaker_state[stage_id]
            state["failure_count"] = 0
            if state["state"] == "half-open":
                logger.info(f"Circuit breaker closed for stage {stage_id}")
                state["state"] = "closed"


# Global recovery engine instance
_recovery_engine: Optional[RecoveryEngine] = None


def get_recovery_engine() -> RecoveryEngine:
    """Get the global recovery engine instance."""
    global _recovery_engine
    if _recovery_engine is None:
        _recovery_engine = RecoveryEngine()
    return _recovery_engine
