"""Exponential backoff with jitter for transient failures."""

from __future__ import annotations

import asyncio
import random
from typing import Any, Callable, Coroutine, Optional, Tuple, Type

from openhedge.exceptions import NetworkError, RateLimitError, ServerError


# Default set of exceptions that should trigger a retry.
RETRYABLE: Tuple[Type[Exception], ...] = (
    NetworkError,
    ServerError,
    RateLimitError,
    ConnectionError,
    TimeoutError,
)


async def with_retry(
    fn: Callable[..., Coroutine[Any, Any, Any]],
    *args: Any,
    max_retries: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 30.0,
    retryable: Optional[Tuple[Type[Exception], ...]] = None,
    **kwargs: Any,
) -> Any:
    """Execute *fn* with exponential backoff on retryable exceptions.

    Parameters
    ----------
    fn:
        Async callable to invoke.
    max_retries:
        Maximum number of retry attempts (0 means execute once).
    base_delay:
        Initial delay in seconds before the first retry.
    max_delay:
        Ceiling for the computed delay.
    retryable:
        Exception types that trigger a retry.  Defaults to common transient
        failures.

    Returns
    -------
    The return value of *fn*.
    """
    retryable = retryable or RETRYABLE
    last_exc: Optional[Exception] = None

    for attempt in range(max_retries + 1):
        try:
            return await fn(*args, **kwargs)
        except retryable as exc:
            last_exc = exc
            if attempt == max_retries:
                raise

            # If we know the server's Retry-After, honour it.
            if isinstance(exc, RateLimitError) and exc.retry_after > 0:
                delay = exc.retry_after
            else:
                delay = min(base_delay * (2 ** attempt), max_delay)
                # Add 0-50 % jitter to avoid thundering herd
                delay += delay * random.random() * 0.5

            await asyncio.sleep(delay)

    # Should be unreachable, but keeps mypy happy.
    raise last_exc  # type: ignore[misc]
