"""Retry logic with exponential backoff."""

from __future__ import annotations

import random
import time
from dataclasses import dataclass
from typing import Callable, TypeVar

T = TypeVar("T")


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

    max_retries: int = 3
    base_delay: float = 1.0  # seconds
    max_delay: float = 60.0  # seconds
    exponential_base: float = 2.0
    jitter: bool = True

    # Exceptions to retry on (if empty, retry on all)
    retry_on: tuple[type[Exception], ...] = ()

    # Exceptions to never retry on
    never_retry_on: tuple[type[Exception], ...] = ()


class RetryError(Exception):
    """Raised when all retries are exhausted."""

    def __init__(self, last_error: Exception, attempts: int):
        self.last_error = last_error
        self.attempts = attempts
        super().__init__(f"Failed after {attempts} attempts: {last_error}")


def should_retry(error: Exception, config: RetryConfig) -> bool:
    """Determine if an error should trigger a retry."""
    # Never retry these
    if config.never_retry_on and isinstance(error, config.never_retry_on):
        return False

    # If retry_on is specified, only retry those
    if config.retry_on:
        return isinstance(error, config.retry_on)

    # Default: retry on most errors
    return True


def calculate_delay(attempt: int, config: RetryConfig) -> float:
    """Calculate delay before next retry with exponential backoff."""
    delay = config.base_delay * (config.exponential_base ** (attempt - 1))
    delay = min(delay, config.max_delay)

    if config.jitter:
        # Add random jitter (0-25% of delay)
        delay = delay * (1 + random.random() * 0.25)

    return delay


def with_retry(
    fn: Callable[[], T],
    config: RetryConfig | None = None,
    on_retry: Callable[[Exception, int], None] | None = None,
) -> T:
    """Execute a function with retry logic.

    Args:
        fn: Function to execute
        config: Retry configuration
        on_retry: Optional callback when a retry occurs

    Returns:
        Result of fn()

    Raises:
        RetryError: If all retries are exhausted
    """
    config = config or RetryConfig()
    last_error: Exception | None = None

    for attempt in range(1, config.max_retries + 2):  # +1 for initial attempt
        try:
            return fn()
        except Exception as e:
            last_error = e

            if attempt > config.max_retries:
                break

            if not should_retry(e, config):
                raise

            delay = calculate_delay(attempt, config)

            if on_retry:
                on_retry(e, attempt)

            time.sleep(delay)

    raise RetryError(last_error, attempt)  # type: ignore


class Retrier:
    """Reusable retry wrapper."""

    def __init__(
        self,
        config: RetryConfig | None = None,
        on_retry: Callable[[Exception, int], None] | None = None,
    ):
        self.config = config or RetryConfig()
        self.on_retry = on_retry

    def __call__(self, fn: Callable[[], T]) -> T:
        """Execute function with retry."""
        return with_retry(fn, self.config, self.on_retry)

    def wrap(self, fn: Callable[..., T]) -> Callable[..., T]:
        """Decorator to wrap a function with retry logic."""

        def wrapped(*args, **kwargs) -> T:
            return with_retry(lambda: fn(*args, **kwargs), self.config, self.on_retry)

        return wrapped
