"""
Tool retry execution with Result-based error handling.

This module provides retry logic specifically designed for tool execution,
where tools return Result objects with is_error flags rather than raising exceptions.

This module is decoupled from app via protocols - RetryPolicy is injected
through the constructor using RetryPolicyProtocol.
"""

import asyncio
import random
from typing import Any, Callable, Dict, Optional

from swarm.config.logging import get_structured_logger
from swarm.config.protocols import RetryPolicyProtocol

logger = get_structured_logger(__name__)


# Mapping from ErrorType to retry_on policy strings
ERROR_TYPE_TO_RETRY_ON: Dict[str, str] = {
    "internal": "internal_error",
    "transient": "transient_error",
    "validation": "validation_error",
    "business": "business_error",
    "user_input": "user_input_error",
}


class ToolRetryExecutor:
    """
    Executes tool functions with Result-based retry logic.

    Unlike WorkflowRetryExecutor which catches exceptions, this executor
    checks Result.is_error and Result.error_type after each execution
    to determine if retry is needed.

    Features:
    - Configurable retry count and backoff strategies
    - Conditional retries based on error types in Result
    - Jitter to prevent thundering herd
    - Detailed retry logging

    Note: Policy is required (no defaults) - follows pure DI pattern.
    Callers must provide a RetryPolicy instance.
    """

    def __init__(self, policy: RetryPolicyProtocol):
        """
        Initialize tool retry executor with policy.

        Args:
            policy: Retry policy configuration (required, no default)
        """
        self.policy = policy
        self._retry_on_errors = (
            set(self.policy.retry_on)
            if self.policy.retry_on
            else {"timeout", "transient_error"}
        )

    def is_retryable(self, error_type: Optional[str]) -> bool:
        """
        Check if error type is retryable based on policy.

        Args:
            error_type: The error type string from Result.error_type.value

        Returns:
            True if this error type should trigger a retry
        """
        if not error_type:
            return False
        retry_on_str = ERROR_TYPE_TO_RETRY_ON.get(error_type, "")
        return retry_on_str in self._retry_on_errors

    def calculate_backoff(self, attempt: int) -> float:
        """
        Calculate backoff delay based on policy.

        Args:
            attempt: Current attempt number (1-based)

        Returns:
            Delay in seconds
        """
        base = float(self.policy.base_seconds)

        match self.policy.backoff:
            case "constant":
                delay = base
            case "linear":
                delay = base * attempt
            case _:  # exponential (default)
                delay = base * (2 ** (attempt - 1))

        if self.policy.jitter:
            delay += delay * 0.25 * random.random()

        return min(delay, 300.0)  # Cap at 5 minutes

    async def execute_with_retry(
        self,
        func: Callable[..., Any],
        *args: Any,
        tool_name: Optional[str] = None,
        **kwargs: Any,
    ) -> Any:
        """
        Execute a tool function with retry logic based on Result errors.

        Args:
            func: Async function that returns a Result object
            *args: Positional arguments for func
            tool_name: Optional tool name for logging
            **kwargs: Keyword arguments for func

        Returns:
            The Result object from the final execution attempt
        """
        result = await func(*args, **kwargs)
        attempt = 1

        while (
            result.is_error
            and self.is_retryable(
                result.error_type.value if result.error_type else None
            )
            and attempt < self.policy.max_retries + 1
        ):
            attempt += 1
            backoff = self.calculate_backoff(attempt - 1)

            logger.info(
                f"Retrying {tool_name or 'tool'} (attempt {attempt}/{self.policy.max_retries + 1})",
                extra={
                    "tool_name": tool_name,
                    "attempt": attempt,
                    "max_attempts": self.policy.max_retries + 1,
                    "error_type": result.error_type.value
                    if result.error_type
                    else None,
                    "backoff_seconds": backoff,
                    "backoff_strategy": self.policy.backoff,
                },
            )

            await asyncio.sleep(backoff)
            result = await func(*args, **kwargs)

        return result
