"""Retry utilities with exponential backoff."""

from __future__ import annotations

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

from src.exceptions import RateLimitError, TranslationError

logger = logging.getLogger(__name__)

T = TypeVar("T")


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

    max_attempts: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    exponential_base: float = 2.0
    jitter: bool = True
    retryable_exceptions: tuple[type[Exception], ...] = field(
        default_factory=lambda: (RateLimitError, ConnectionError, TimeoutError)
    )

    def get_delay(self, attempt: int) -> float:
        """Calculate delay for the given attempt number.

        Uses exponential backoff with optional jitter.

        Args:
            attempt: The attempt number (0-indexed).

        Returns:
            Delay in seconds before the next retry.
        """
        delay = self.base_delay * (self.exponential_base**attempt)
        delay = min(delay, self.max_delay)

        if self.jitter:
            # Add random jitter of +/- 25%
            jitter_range = delay * 0.25
            delay += random.uniform(-jitter_range, jitter_range)

        return max(0, delay)


def retry_with_backoff(
    config: RetryConfig | None = None,
    on_retry: Callable[[Exception, int], None] | None = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
    """Decorator for retrying functions with exponential backoff.

    Args:
        config: Retry configuration. Uses defaults if not provided.
        on_retry: Optional callback called before each retry with (exception, attempt).

    Returns:
        Decorator function.
    """
    if config is None:
        config = RetryConfig()

    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> T:
            last_exception: Exception | None = None

            for attempt in range(config.max_attempts):
                try:
                    return func(*args, **kwargs)
                except config.retryable_exceptions as e:
                    last_exception = e

                    # Don't sleep after the last attempt
                    if attempt < config.max_attempts - 1:
                        # Check for rate limit specific retry-after
                        if isinstance(e, RateLimitError) and e.retry_after:
                            delay = e.retry_after
                        else:
                            delay = config.get_delay(attempt)

                        logger.warning(
                            "Attempt %d/%d failed: %s. Retrying in %.1fs...",
                            attempt + 1,
                            config.max_attempts,
                            str(e),
                            delay,
                        )

                        if on_retry:
                            on_retry(e, attempt)

                        time.sleep(delay)
                except Exception:
                    # Non-retryable exception, re-raise immediately
                    raise

            # All attempts exhausted
            if last_exception:
                raise last_exception
            raise TranslationError("All retry attempts failed")

        return wrapper

    return decorator


def async_retry_with_backoff(
    config: RetryConfig | None = None,
    on_retry: Callable[[Exception, int], None] | None = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
    """Async decorator for retrying functions with exponential backoff.

    Args:
        config: Retry configuration. Uses defaults if not provided.
        on_retry: Optional callback called before each retry with (exception, attempt).

    Returns:
        Decorator function for async functions.
    """
    if config is None:
        config = RetryConfig()

    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Coroutine[Any, Any, T]]:
        @functools.wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> T:
            last_exception: Exception | None = None

            for attempt in range(config.max_attempts):
                try:
                    return await func(*args, **kwargs)
                except config.retryable_exceptions as e:
                    last_exception = e

                    # Don't sleep after the last attempt
                    if attempt < config.max_attempts - 1:
                        # Check for rate limit specific retry-after
                        if isinstance(e, RateLimitError) and e.retry_after:
                            delay = e.retry_after
                        else:
                            delay = config.get_delay(attempt)

                        logger.warning(
                            "Attempt %d/%d failed: %s. Retrying in %.1fs...",
                            attempt + 1,
                            config.max_attempts,
                            str(e),
                            delay,
                        )

                        if on_retry:
                            on_retry(e, attempt)

                        await asyncio.sleep(delay)
                except Exception:
                    # Non-retryable exception, re-raise immediately
                    raise

            # All attempts exhausted
            if last_exception:
                raise last_exception
            raise TranslationError("All retry attempts failed")

        return wrapper

    return decorator


class RetryState:
    """Tracks retry state for manual retry loops."""

    def __init__(self, config: RetryConfig | None = None) -> None:
        """Initialize retry state.

        Args:
            config: Retry configuration.
        """
        self.config = config or RetryConfig()
        self.attempt = 0
        self.last_exception: Exception | None = None

    def should_retry(self, exception: Exception) -> bool:
        """Check if we should retry after an exception.

        Args:
            exception: The exception that was raised.

        Returns:
            True if we should retry, False otherwise.
        """
        self.last_exception = exception

        if not isinstance(exception, self.config.retryable_exceptions):
            return False

        if self.attempt >= self.config.max_attempts - 1:
            return False

        return True

    def wait(self) -> None:
        """Wait before the next retry attempt."""
        if isinstance(self.last_exception, RateLimitError):
            delay = self.last_exception.retry_after or self.config.get_delay(self.attempt)
        else:
            delay = self.config.get_delay(self.attempt)

        time.sleep(delay)
        self.attempt += 1

    async def async_wait(self) -> None:
        """Async wait before the next retry attempt."""
        if isinstance(self.last_exception, RateLimitError):
            delay = self.last_exception.retry_after or self.config.get_delay(self.attempt)
        else:
            delay = self.config.get_delay(self.attempt)

        await asyncio.sleep(delay)
        self.attempt += 1

    def reset(self) -> None:
        """Reset retry state for a new operation."""
        self.attempt = 0
        self.last_exception = None