from __future__ import annotations

import asyncio
import os
import random
from typing import Awaitable, Callable, Optional


class RetryError(Exception):
    pass


def _get_int_env(name: str, default: int) -> int:
    try:
        return int(os.getenv(name, str(default)))
    except Exception:
        return default


def _get_float_env(name: str, default: float) -> float:
    try:
        return float(os.getenv(name, str(default)))
    except Exception:
        return default


def is_retryable_error(exc: BaseException) -> bool:
    msg = str(exc).lower()
    retry_terms = [
        "rate limit",
        "quota",
        "timeout",
        "timed out",
        "temporarily",
        "unavailable",
        "service unavailable",
        "connection reset",
        "connection aborted",
        "429",
        "502",
        "503",
        "504",
    ]
    return any(term in msg for term in retry_terms)


async def retry_async(
    op_factory: Callable[[], Awaitable],
    *,
    max_attempts: Optional[int] = None,
    base_delay_s: Optional[float] = None,
    max_delay_s: Optional[float] = None,
    jitter: bool = True,
    is_retryable: Callable[[BaseException], bool] = is_retryable_error,
):
    """Retry an async operation with exponential backoff.

    op_factory should return a fresh awaitable each attempt.
    """
    attempts = max_attempts or _get_int_env("LLMRING_RETRY_ATTEMPTS", 2)
    base = (
        base_delay_s
        if base_delay_s is not None
        else _get_float_env("LLMRING_RETRY_BASE_S", 0.25)
    )
    max_delay = (
        max_delay_s
        if max_delay_s is not None
        else _get_float_env("LLMRING_RETRY_MAX_S", 2.0)
    )

    last_exc: Optional[BaseException] = None
    for attempt in range(1, attempts + 1):
        try:
            return await op_factory()
        except Exception as exc:  # noqa: BLE001
            last_exc = exc
            if attempt >= attempts or not is_retryable(exc):
                break
            # exponential backoff with jitter
            delay = min(max_delay, base * (2 ** (attempt - 1)))
            if jitter:
                delay = delay * (0.5 + random.random() * 0.5)
            await asyncio.sleep(delay)

    raise RetryError(str(last_exc) if last_exc else "unknown error") from last_exc
