"""Retry helpers.

This module provides `async_retry`, an exponential backoff retry wrapper with
optional jitter. It also attempts to honor `Retry-After` hints for rate limits
when available on raised exceptions.
"""

from __future__ import annotations

import asyncio
import random
from typing import Awaitable, Callable, Iterable, TypeVar

T = TypeVar("T")


async def async_retry(
    func: Callable[[], Awaitable[T]],
    max_attempts: int = 3,
    backoff_base: float = 2.0,
    jitter: bool = True,
    retry_on: Iterable[type[BaseException]] = (),
) -> T:
    """Retry an async callable using exponential backoff.

    Args:
        func: A zero-arg async callable.
        max_attempts: Maximum number of attempts (including the first).
        backoff_base: Base for exponential backoff.
        jitter: If True, applies small random jitter to reduce thundering herd.
        retry_on: Exception types to retry on.
    """

    retry_on_tuple = tuple(retry_on)
    if max_attempts < 1:
        raise ValueError("max_attempts must be >= 1")

    last_exc: BaseException | None = None
    for attempt in range(max_attempts):
        try:
            return await func()
        except Exception as e:  # noqa: BLE001
            last_exc = e
            if retry_on_tuple and not isinstance(e, retry_on_tuple):
                raise
            if attempt >= max_attempts - 1:
                raise

            retry_after = _extract_retry_after_seconds(e)
            if retry_after is None:
                # Exponential backoff: base^attempt.
                retry_after = float(backoff_base**attempt)

            if jitter:
                retry_after = retry_after * random.uniform(0.6, 1.4)

            await asyncio.sleep(max(0.0, retry_after))

    assert last_exc is not None
    raise last_exc


def _extract_retry_after_seconds(exc: BaseException) -> float | None:
    """Best-effort extraction of a Retry-After delay from common error shapes."""
    # Common: adapters attach retry_after_seconds.
    for attr in ("retry_after_seconds", "retry_after"):
        if hasattr(exc, attr):
            val = getattr(exc, attr)
            if isinstance(val, (int, float)):
                return float(val)
            if isinstance(val, str):
                try:
                    return float(val)
                except ValueError:
                    pass

    # httpx.HTTPStatusError has .response.headers
    if hasattr(exc, "response"):
        response = getattr(exc, "response", None)
        headers = getattr(response, "headers", None)
        if headers:
            retry_after = headers.get("Retry-After")
            if retry_after:
                try:
                    return float(retry_after)
                except ValueError:
                    return None
    return None

