"""
Retry logic for agent operations.

Provides configurable retry policies with exponential backoff and jitter
for handling transient failures in agent execution.
"""

import asyncio
import random
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any, TypeVar

from persuade.agents.base import AgentPhaseResult, TwoPhaseAgent
from persuade.logging import get_logger

logger = get_logger()

T = TypeVar("T")


@dataclass
class RetryPolicy:
    """Configuration for retry behavior.

    Implements exponential backoff with optional jitter to prevent
    thundering herd problems in distributed systems.

    Attributes:
        max_attempts: Maximum number of attempts (including initial).
                     Set to 1 to disable retries.
        initial_delay_seconds: Delay before first retry.
        backoff_multiplier: Multiply delay by this after each attempt.
        max_delay_seconds: Cap on delay between retries.
        jitter: If True, add random jitter (50-150% of delay).
        retryable_errors: Exception types that should trigger retry.

    Example:
        # Retry up to 3 times with exponential backoff
        policy = RetryPolicy(
            max_attempts=3,
            initial_delay_seconds=1.0,
            backoff_multiplier=2.0,
        )
        # Delays: 1s, 2s, 4s (with jitter: 0.5-1.5s, 1-3s, 2-6s)
    """

    max_attempts: int = 3
    initial_delay_seconds: float = 1.0
    backoff_multiplier: float = 2.0
    max_delay_seconds: float = 60.0
    jitter: bool = True
    retryable_errors: tuple[type[Exception], ...] = field(
        default_factory=lambda: (TimeoutError, ConnectionError)
    )


# Preset policies for common use cases
RETRY_NONE = RetryPolicy(max_attempts=1)
"""No retry - fail immediately on first error."""

RETRY_QUICK = RetryPolicy(max_attempts=2, initial_delay_seconds=0.5)
"""Quick retry - one retry after 0.5s for fast operations."""

RETRY_STANDARD = RetryPolicy(max_attempts=3, initial_delay_seconds=1.0)
"""Standard retry - three attempts with 1s base delay."""

RETRY_PERSISTENT = RetryPolicy(
    max_attempts=5,
    initial_delay_seconds=2.0,
    max_delay_seconds=30.0,
)
"""Persistent retry - five attempts for critical operations."""


def calculate_delay(policy: RetryPolicy, attempt: int) -> float:
    """Calculate delay before next retry with optional jitter.

    Args:
        policy: The retry policy configuration
        attempt: The attempt number (0-indexed)

    Returns:
        Delay in seconds before next retry
    """
    delay = min(
        policy.initial_delay_seconds * (policy.backoff_multiplier**attempt),
        policy.max_delay_seconds,
    )
    if policy.jitter:
        # Add 50-150% jitter to prevent thundering herd
        delay *= 0.5 + random.random()
    return delay


async def with_retry(
    operation: Callable[[], Awaitable[T]],
    policy: RetryPolicy,
    operation_name: str = "operation",
) -> T:
    """Execute an async operation with retry logic.

    Args:
        operation: Async callable to execute
        policy: Retry configuration
        operation_name: Name for logging (helps with debugging)

    Returns:
        Result of successful operation

    Raises:
        Last exception if all attempts fail

    Example:
        result = await with_retry(
            lambda: fetch_data(),
            RETRY_STANDARD,
            "fetch_data",
        )
    """
    last_error: Exception | None = None

    for attempt in range(policy.max_attempts):
        try:
            return await operation()
        except policy.retryable_errors as e:
            last_error = e
            if attempt < policy.max_attempts - 1:
                delay = calculate_delay(policy, attempt)
                logger.warning(
                    f"Retry {attempt + 1}/{policy.max_attempts} for {operation_name}",
                    error=str(e),
                    delay_seconds=round(delay, 2),
                )
                await asyncio.sleep(delay)
            else:
                logger.error(
                    f"All {policy.max_attempts} attempts failed for {operation_name}",
                    error=str(e),
                )

    # Should not reach here, but type checker needs this
    if last_error is not None:
        raise last_error
    raise RuntimeError(f"Unexpected state in retry loop for {operation_name}")


class RetryableAgent(TwoPhaseAgent):
    """Wrapper that adds retry logic to any TwoPhaseAgent.

    Wraps an existing agent and adds configurable retry behavior
    for both plan and execute phases. Different policies can be
    used for each phase.

    Example:
        # Wrap an existing agent with retry logic
        agent = RetryableAgent(
            IteratePRAgent(pr_number=67),
            plan_policy=RETRY_STANDARD,
            execute_policy=RETRY_QUICK,
        )
        result = await agent.run(auto_approve=True)

    Note:
        The execute phase typically uses fewer retries because
        partial execution may have already occurred.
    """

    def __init__(
        self,
        agent: TwoPhaseAgent,
        plan_policy: RetryPolicy | None = None,
        execute_policy: RetryPolicy | None = None,
    ):
        """Initialize the retryable wrapper.

        Args:
            agent: The agent to wrap with retry logic
            plan_policy: Retry policy for plan phase. Defaults to RETRY_STANDARD.
            execute_policy: Retry policy for execute phase. Defaults to RETRY_QUICK.
        """
        self.agent = agent
        self.plan_policy = plan_policy or RETRY_STANDARD
        self.execute_policy = execute_policy or RETRY_QUICK

    async def plan(self, **kwargs: Any) -> AgentPhaseResult:
        """Execute plan phase with retry.

        Retries the entire plan phase on retryable errors.
        Non-retryable errors or AgentPhaseResult failures are not retried.
        """
        try:
            return await with_retry(
                lambda: self.agent.plan(**kwargs),
                self.plan_policy,
                f"{self.agent.__class__.__name__}.plan",
            )
        except Exception as e:
            return AgentPhaseResult(
                success=False,
                error=f"Plan failed after {self.plan_policy.max_attempts} attempts: {e}",
            )

    async def execute(self, plan_result: AgentPhaseResult, **kwargs: Any) -> AgentPhaseResult:
        """Execute phase with retry.

        Retries the entire execute phase on retryable errors.
        Uses a more conservative policy by default since partial
        execution may have occurred.
        """
        try:
            return await with_retry(
                lambda: self.agent.execute(plan_result, **kwargs),
                self.execute_policy,
                f"{self.agent.__class__.__name__}.execute",
            )
        except Exception as e:
            return AgentPhaseResult(
                success=False,
                error=f"Execute failed after {self.execute_policy.max_attempts} attempts: {e}",
            )
