"""
Retry and fallback logic for agent execution.

This module provides retry mechanisms with exponential backoff and
fallback agent support using the Tenacity library. It ensures robust
execution even when agents encounter transient failures.

Extended with graceful degradation support for partial results when
all agents fail (EXEC-13, EXEC-14, EXEC-15).

Requirements: AGENT-08 (retry and fallback), EXEC-13, EXEC-14, EXEC-15
"""

import logging
from typing import Callable, List, Optional, Any, Tuple, TYPE_CHECKING

from tenacity import (
    retry,
    stop_after_attempt,
    stop_after_delay,
    wait_exponential,
    retry_if_exception_type,
    before_sleep_log,
)

if TYPE_CHECKING:
    from gsd_rlm.execution.degradation import GracefulDegradation, PartialResult

logger = logging.getLogger(__name__)


class AgentExecutionError(Exception):
    """Raised when agent execution fails."""

    pass


class FallbackExhaustedError(Exception):
    """All fallback agents failed."""

    pass


class RetryConfig:
    """Configuration for retry behavior.

    This class provides a simple way to configure retry behavior
    without dealing with Tenacity's decorator parameters directly.

    Attributes:
        max_attempts: Maximum number of retry attempts
        max_delay: Maximum total delay in seconds
        min_wait: Minimum wait between retries in seconds
        max_wait: Maximum wait between retries in seconds
    """

    def __init__(
        self,
        max_attempts: int = 3,
        max_delay: float = 60.0,
        min_wait: float = 2.0,
        max_wait: float = 30.0,
    ):
        """Initialize retry configuration.

        Args:
            max_attempts: Maximum retry attempts (default: 3)
            max_delay: Maximum total delay in seconds (default: 60)
            min_wait: Minimum wait between retries (default: 2s)
            max_wait: Maximum wait between retries (default: 30s)
        """
        self.max_attempts = max_attempts
        self.max_delay = max_delay
        self.min_wait = min_wait
        self.max_wait = max_wait


class RetryableAgent:
    """Agent wrapper with retry and fallback logic (AGENT-08).

    This class wraps an agent with automatic retry behavior and
    optional fallback agents. When the primary agent fails, it
    automatically retries with exponential backoff. If all retries
    fail, it tries fallback agents in sequence.

    Extended with graceful degradation support (EXEC-13, EXEC-14, EXEC-15):
    When degradation is configured, failures return partial results instead
    of raising FallbackExhaustedError.

    Example:
        ```python
        from gsd_rlm.execution.retry import RetryableAgent, RetryConfig

        # Create retryable agent with fallbacks
        retryable = RetryableAgent(
            primary_agent=main_agent,
            fallback_agents=[backup_agent1, backup_agent2],
            config=RetryConfig(max_attempts=3, max_delay=60)
        )

        # Execute with automatic retry and fallback
        result = retryable.execute(message)
        ```
    """

    def __init__(
        self,
        primary_agent: Any,
        fallback_agents: Optional[List[Any]] = None,
        config: Optional[RetryConfig] = None,
        degradation: Optional["GracefulDegradation"] = None,
        on_partial_success: Optional[Callable[["PartialResult"], None]] = None,
    ):
        """Initialize the retryable agent wrapper.

        Args:
            primary_agent: The main agent to execute first
            fallback_agents: List of fallback agents to try if primary fails
            config: Retry configuration (uses defaults if None)
            degradation: Optional GracefulDegradation instance for partial results
            on_partial_success: Callback when partial success occurs (with degradation)
        """
        self.primary_agent = primary_agent
        self.fallback_agents = fallback_agents or []
        self.config = config or RetryConfig()
        self.degradation = degradation
        self.on_partial_success = on_partial_success

    def execute(self, message: Any) -> Any:
        """Execute with retry, then fallback agents.

        This method first tries the primary agent with exponential backoff
        retry. If all retries fail, it tries each fallback agent in sequence.

        If degradation is configured and all agents fail, returns None
        instead of raising FallbackExhaustedError. Use execute_with_degradation()
        for explicit degradation handling.

        Args:
            message: The input message to process

        Returns:
            The result from the first successful agent, or None if all failed
            and degradation is configured

        Raises:
            FallbackExhaustedError: If all agents fail and no degradation configured
        """
        errors = []

        # Try primary agent with retry
        try:
            return self._execute_with_retry(self.primary_agent, message)
        except AgentExecutionError as e:
            errors.append(f"Primary agent failed: {e}")
            logger.error(errors[-1])

        # Try fallback agents
        for i, fallback in enumerate(self.fallback_agents):
            try:
                agent_name = getattr(fallback, "name", f"fallback-{i}")
                logger.info(f"Trying fallback agent {i + 1}: {agent_name}")
                return self._execute_with_retry(fallback, message)
            except AgentExecutionError as e:
                errors.append(
                    f"Fallback {i + 1} ({getattr(fallback, 'name', i)}) failed: {e}"
                )
                logger.error(errors[-1])

        # All failed - handle with degradation or raise
        if self.degradation is not None:
            logger.warning("All agents failed, returning None with degradation enabled")
            return None

        raise FallbackExhaustedError("All agents failed:\n" + "\n".join(errors))

    def _execute_with_retry(self, agent: Any, message: Any) -> Any:
        """Execute with exponential backoff retry.

        This method is decorated with Tenacity's retry decorator for
        automatic retry with exponential backoff.

        Args:
            agent: The agent to execute
            message: The input message

        Returns:
            The agent's result

        Raises:
            AgentExecutionError: If execution fails after all retries
        """

        @retry(
            stop=(
                stop_after_attempt(self.config.max_attempts)
                | stop_after_delay(self.config.max_delay)
            ),
            wait=wait_exponential(
                multiplier=1,
                min=self.config.min_wait,
                max=self.config.max_wait,
            ),
            retry=retry_if_exception_type(AgentExecutionError),
            before_sleep=before_sleep_log(logger, logging.WARNING),
            reraise=True,
        )
        def _do_retry():
            try:
                result = agent.process(message)

                if not result:
                    raise AgentExecutionError(
                        f"Agent {getattr(agent, 'name', agent)} returned no result"
                    )

                return result[0] if isinstance(result, list) else result

            except AgentExecutionError:
                raise
            except Exception as e:
                logger.warning(f"Agent {getattr(agent, 'name', agent)} failed: {e}")
                raise AgentExecutionError(str(e)) from e

        return _do_retry()

    def execute_with_degradation(
        self, message: Any
    ) -> Tuple[Any, Optional["PartialResult"]]:
        """Execute with explicit degradation handling.

        This method provides explicit access to partial results when
        degradation is configured. Unlike execute(), it returns a tuple
        with the result and an optional PartialResult.

        Args:
            message: The input message to process

        Returns:
            Tuple of (result, partial_result):
            - On success: (result, None)
            - On failure with degradation: (None, PartialResult)
            - On failure without degradation: raises FallbackExhaustedError

        Raises:
            FallbackExhaustedError: If all agents fail and no degradation configured

        Example:
            ```python
            retryable = RetryableAgent(
                primary_agent=agent,
                degradation=GracefulDegradation(min_success_rate=0.5),
                on_partial_success=lambda p: print(f"Partial: {p.success_rate}")
            )

            result, partial = retryable.execute_with_degradation(message)
            if partial is not None:
                # Handle partial failure
                print(f"Got {partial.completed_tasks}/{partial.total_tasks}")
            else:
                # Full success
                print(f"Result: {result}")
            ```
        """
        from gsd_rlm.execution.degradation import PartialResult, TaskStatus
        from dataclasses import dataclass

        @dataclass
        class TaskResult:
            task_id: str
            content: Any
            success: bool
            error: str = None

        errors = []
        task_results = []

        # Try primary agent with retry
        try:
            result = self._execute_with_retry(self.primary_agent, message)
            return (result, None)
        except AgentExecutionError as e:
            errors.append(f"Primary agent failed: {e}")
            task_results.append(TaskResult("primary", None, False, str(e)))
            logger.error(errors[-1])

        # Try fallback agents
        for i, fallback in enumerate(self.fallback_agents):
            try:
                agent_name = getattr(fallback, "name", f"fallback-{i}")
                logger.info(f"Trying fallback agent {i + 1}: {agent_name}")
                result = self._execute_with_retry(fallback, message)
                return (result, None)
            except AgentExecutionError as e:
                task_id = f"fallback-{i}"
                errors.append(
                    f"Fallback {i + 1} ({getattr(fallback, 'name', i)}) failed: {e}"
                )
                task_results.append(TaskResult(task_id, None, False, str(e)))
                logger.error(errors[-1])

        # All failed - handle with degradation
        if self.degradation is not None:
            partial = self.degradation.collect_partial_results(
                wave_id=0, task_results=task_results
            )

            # Call callback if provided
            if self.on_partial_success is not None:
                try:
                    self.on_partial_success(partial)
                except Exception as e:
                    logger.warning(f"on_partial_success callback failed: {e}")

            return (None, partial)

        raise FallbackExhaustedError("All agents failed:\n" + "\n".join(errors))


def with_retry(
    max_attempts: int = 3,
    max_delay: float = 60.0,
    min_wait: float = 2.0,
    max_wait: float = 30.0,
):
    """Decorator to add retry to any function.

    This is a convenience decorator that adds retry behavior to any
    function without requiring the RetryableAgent wrapper.

    Args:
        max_attempts: Maximum retry attempts (default: 3)
        max_delay: Maximum total delay in seconds (default: 60)
        min_wait: Minimum wait between retries (default: 2s)
        max_wait: Maximum wait between retries (default: 30s)

    Returns:
        A decorator that adds retry behavior

    Example:
        ```python
        from gsd_rlm.execution.retry import with_retry

        @with_retry(max_attempts=5, max_delay=120)
        def call_external_api(data):
            # This will automatically retry on failure
            return requests.post(url, json=data)
        ```
    """
    return retry(
        stop=(stop_after_attempt(max_attempts) | stop_after_delay(max_delay)),
        wait=wait_exponential(multiplier=1, min=min_wait, max=max_wait),
        reraise=True,
    )
