"""
Recovery Engine - Retry, Backoff, and Circuit Breaker Logic.

Implements multiple recovery strategies for resilient pipeline execution.
"""

from __future__ import annotations

import asyncio
import logging
import random
import time
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Awaitable, Callable

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


class RecoveryStrategy(str, Enum):
    """Available recovery strategies."""

    IMMEDIATE_RETRY = "immediate_retry"  # Retry immediately
    EXPONENTIAL_BACKOFF = "exponential_backoff"  # Retry with exponential delays
    FIXED_BACKOFF = "fixed_backoff"  # Retry with fixed delays
    CIRCUIT_BREAKER = "circuit_breaker"  # Stop retrying after threshold
    FALLBACK = "fallback"  # Use alternative handler
    SKIP = "skip"  # Skip and continue


class FailureType(str, Enum):
    """Classification of failure types."""

    TRANSIENT = "transient"  # Network timeout, temp unavailable - retry
    PERMANENT = "permanent"  # Invalid input, auth failure - don't retry
    RESOURCE = "resource"  # Out of memory, disk full - wait and retry
    EXTERNAL = "external"  # External service failure - circuit breaker
    UNKNOWN = "unknown"  # Unclassified error


class RecoveryConfig(BaseModel):
    """Configuration for recovery behavior."""

    strategy: RecoveryStrategy = RecoveryStrategy.EXPONENTIAL_BACKOFF
    max_retries: int = Field(default=3, ge=0, le=10)
    base_delay: float = Field(default=1.0, ge=0.0)
    max_delay: float = Field(default=60.0, ge=0.0)
    exponential_base: float = Field(default=2.0, ge=1.0, le=4.0)
    jitter: bool = True  # Add randomness to delays

    # Circuit breaker settings
    failure_threshold: int = Field(default=5, ge=1)
    reset_timeout: float = Field(default=60.0, ge=0.0)

    # Fallback handler (fully qualified function name)
    fallback_handler: str | None = None

    # LLM-assisted diagnosis
    enable_llm_diagnosis: bool = True

    model_config = {"frozen": True}


class RecoveryAttempt(BaseModel):
    """Record of a recovery attempt."""

    attempt_number: int
    strategy: RecoveryStrategy
    success: bool
    error: str | None = None
    delay_used: float = 0.0
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

    # Diagnosis
    failure_type: FailureType = FailureType.UNKNOWN
    llm_diagnosis: str | None = None
    suggested_fix: str | None = None


class RecoveryResult(BaseModel):
    """Result of recovery process."""

    success: bool
    final_error: str | None = None
    attempts: list[RecoveryAttempt] = Field(default_factory=list)
    data: dict[str, Any] = Field(default_factory=dict)
    used_fallback: bool = False
    circuit_opened: bool = False


class CircuitState(str, Enum):
    """Circuit breaker states."""

    CLOSED = "closed"  # Normal operation
    OPEN = "open"  # Blocking calls
    HALF_OPEN = "half_open"  # Testing recovery


class CircuitBreaker:
    """
    Circuit breaker pattern implementation.

    Prevents cascade failures by stopping calls after threshold failures.
    """

    def __init__(
        self,
        failure_threshold: int = 5,
        reset_timeout: float = 60.0,
    ):
        self.failure_threshold = failure_threshold
        self.reset_timeout = reset_timeout
        self._failure_count = 0
        self._state = CircuitState.CLOSED
        self._last_failure_time: float | None = None

    @property
    def state(self) -> CircuitState:
        """Get current circuit state, considering timeout."""
        if self._state == CircuitState.OPEN:
            if self._last_failure_time is not None:
                elapsed = time.monotonic() - self._last_failure_time
                if elapsed >= self.reset_timeout:
                    self._state = CircuitState.HALF_OPEN
        return self._state

    def record_success(self) -> None:
        """Record a successful call."""
        self._failure_count = 0
        self._state = CircuitState.CLOSED

    def record_failure(self) -> None:
        """Record a failed call."""
        self._failure_count += 1
        self._last_failure_time = time.monotonic()
        if self._failure_count >= self.failure_threshold:
            self._state = CircuitState.OPEN

    def can_execute(self) -> bool:
        """Check if execution is allowed."""
        state = self.state  # Triggers timeout check
        return state != CircuitState.OPEN

    def reset(self) -> None:
        """Reset the circuit breaker."""
        self._failure_count = 0
        self._state = CircuitState.CLOSED
        self._last_failure_time = None


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

    # Patterns for transient errors
    TRANSIENT_PATTERNS = [
        "timeout",
        "connection refused",
        "connection reset",
        "temporarily unavailable",
        "rate limit",
        "503",
        "502",
        "504",
        "EAGAIN",
        "ECONNRESET",
    ]

    # Patterns for permanent errors
    PERMANENT_PATTERNS = [
        "invalid",
        "unauthorized",
        "forbidden",
        "not found",
        "400",
        "401",
        "403",
        "404",
        "validation error",
        "syntax error",
    ]

    # Patterns for resource errors
    RESOURCE_PATTERNS = [
        "out of memory",
        "disk full",
        "no space left",
        "quota exceeded",
        "resource exhausted",
        "ENOMEM",
        "ENOSPC",
    ]

    @classmethod
    def classify(cls, error: Exception | str) -> FailureType:
        """Classify an error by type."""
        error_str = str(error).lower()

        for pattern in cls.TRANSIENT_PATTERNS:
            if pattern in error_str:
                return FailureType.TRANSIENT

        for pattern in cls.PERMANENT_PATTERNS:
            if pattern in error_str:
                return FailureType.PERMANENT

        for pattern in cls.RESOURCE_PATTERNS:
            if pattern in error_str:
                return FailureType.RESOURCE

        return FailureType.UNKNOWN


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

    Manages retry logic, backoff, circuit breakers, and fallbacks.
    """

    def __init__(
        self,
        config: RecoveryConfig | None = None,
        circuit_breakers: dict[str, CircuitBreaker] | None = None,
    ):
        self.config = config or RecoveryConfig()
        self._circuit_breakers = circuit_breakers or {}
        self._classifier = ErrorClassifier()

    def get_circuit_breaker(self, key: str) -> CircuitBreaker:
        """Get or create a circuit breaker for a key (e.g., stage ID)."""
        if key not in self._circuit_breakers:
            self._circuit_breakers[key] = CircuitBreaker(
                failure_threshold=self.config.failure_threshold,
                reset_timeout=self.config.reset_timeout,
            )
        return self._circuit_breakers[key]

    def calculate_delay(self, attempt: int) -> float:
        """Calculate delay for a retry attempt."""
        if self.config.strategy == RecoveryStrategy.IMMEDIATE_RETRY:
            return 0.0

        if self.config.strategy == RecoveryStrategy.FIXED_BACKOFF:
            delay = self.config.base_delay
        else:  # Exponential backoff
            delay = self.config.base_delay * (self.config.exponential_base ** attempt)

        # Cap at max delay
        delay = min(delay, self.config.max_delay)

        # Add jitter
        if self.config.jitter:
            jitter_range = delay * 0.2  # ±20%
            delay += random.uniform(-jitter_range, jitter_range)

        return max(0.0, delay)

    async def execute_with_recovery(
        self,
        handler: Callable[[], Awaitable[dict[str, Any]]],
        circuit_key: str | None = None,
        fallback: Callable[[], Awaitable[dict[str, Any]]] | None = None,
    ) -> RecoveryResult:
        """
        Execute a handler with recovery logic.

        Args:
            handler: The async function to execute
            circuit_key: Key for circuit breaker (e.g., stage ID)
            fallback: Optional fallback handler

        Returns:
            RecoveryResult with success status and attempts
        """
        attempts: list[RecoveryAttempt] = []
        circuit_breaker = self.get_circuit_breaker(circuit_key) if circuit_key else None

        # Check circuit breaker
        if circuit_breaker and not circuit_breaker.can_execute():
            return RecoveryResult(
                success=False,
                final_error="Circuit breaker is open - too many recent failures",
                circuit_opened=True,
            )

        last_error: str | None = None

        for attempt in range(self.config.max_retries + 1):
            try:
                result = await handler()

                # Success
                if circuit_breaker:
                    circuit_breaker.record_success()

                attempts.append(RecoveryAttempt(
                    attempt_number=attempt,
                    strategy=self.config.strategy,
                    success=True,
                ))

                return RecoveryResult(
                    success=True,
                    attempts=attempts,
                    data=result,
                )

            except Exception as e:
                last_error = str(e)
                failure_type = ErrorClassifier.classify(e)

                if circuit_breaker:
                    circuit_breaker.record_failure()

                attempt_record = RecoveryAttempt(
                    attempt_number=attempt,
                    strategy=self.config.strategy,
                    success=False,
                    error=last_error,
                    failure_type=failure_type,
                )
                attempts.append(attempt_record)

                logger.warning(
                    f"Recovery attempt {attempt + 1}/{self.config.max_retries + 1} failed: {last_error}"
                )

                # Don't retry permanent errors
                if failure_type == FailureType.PERMANENT:
                    break

                # Calculate delay for next attempt
                if attempt < self.config.max_retries:
                    delay = self.calculate_delay(attempt)
                    attempt_record.delay_used = delay
                    if delay > 0:
                        await asyncio.sleep(delay)

        # All retries exhausted - try fallback
        if fallback:
            try:
                result = await fallback()
                return RecoveryResult(
                    success=True,
                    attempts=attempts,
                    data=result,
                    used_fallback=True,
                )
            except Exception as e:
                last_error = f"Fallback also failed: {e}"

        return RecoveryResult(
            success=False,
            final_error=last_error,
            attempts=attempts,
            circuit_opened=circuit_breaker is not None and not circuit_breaker.can_execute(),
        )

    def should_retry(self, failure_type: FailureType, attempt: int) -> bool:
        """Determine if a retry should be attempted."""
        if attempt >= self.config.max_retries:
            return False

        if failure_type == FailureType.PERMANENT:
            return False

        if self.config.strategy == RecoveryStrategy.SKIP:
            return False

        return True

    def reset_circuit_breaker(self, key: str) -> None:
        """Reset a circuit breaker."""
        if key in self._circuit_breakers:
            self._circuit_breakers[key].reset()

    def get_stats(self) -> dict[str, Any]:
        """Get recovery engine statistics."""
        circuit_stats = {}
        for key, cb in self._circuit_breakers.items():
            circuit_stats[key] = {
                "state": cb.state.value,
                "failure_count": cb._failure_count,
            }

        return {
            "strategy": self.config.strategy.value,
            "max_retries": self.config.max_retries,
            "circuit_breakers": circuit_stats,
        }
