from __future__ import annotations

import asyncio
import random
from enum import Enum
from typing import Any, TypeVar

from pydantic import BaseModel, Field

from nexus.resilience.errors import ProviderError, ToolFailure

T = TypeVar("T")


class RetryStrategy(str, Enum):
    FIXED_PROMPT = "fixed_prompt"
    VARIED_PROMPT = "varied_prompt"
    ROTATE_TOOL = "rotate_tool"
    ESCALATE = "escalate"


class RetryConfig(BaseModel):
    max_attempts: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    multiplier: float = 2.0
    jitter: float = 0.1
    strategy: RetryStrategy = RetryStrategy.FIXED_PROMPT
    retry_on: list[type] = Field(default_factory=lambda: [ProviderError, ToolFailure])


class SmartRetry:
    """Exponential backoff retry helper with jitter."""

    def __init__(self, config: RetryConfig | None = None) -> None:
        self.config = config or RetryConfig()

    async def execute(self, fn, *args, **kwargs) -> T:
        """Run an async callable with strategy-aware retry behavior."""
        last_exc: Exception | None = None
        for attempt in range(self.config.max_attempts):
            try:
                return await fn(*args, **kwargs)
            except Exception as exc:
                last_exc = exc
                if not any(isinstance(exc, t) for t in self.config.retry_on):
                    raise
                if attempt >= self.config.max_attempts - 1:
                    break
                if hasattr(exc, "context") and isinstance(getattr(exc, "context"), dict):
                    exc.context["attempt"] = attempt + 1
                    exc.context["strategy"] = self.config.strategy.value
                await asyncio.sleep(self._compute_delay(attempt))

        assert last_exc is not None
        if hasattr(last_exc, "context") and isinstance(getattr(last_exc, "context"), dict):
            last_exc.context["max_attempts"] = self.config.max_attempts
        raise last_exc

    def _compute_delay(self, attempt: int) -> float:
        """Compute capped exponential delay with additive random jitter."""
        base = min(self.config.base_delay * (self.config.multiplier ** attempt), self.config.max_delay)
        jitter_range = self.config.jitter * base
        return max(0.0, base + random.uniform(-jitter_range, jitter_range))
