"""Async retry helper used by every pipeline stage.

Design:
  - Exponential backoff with jitter (2s, 4s, 8s base).
  - Retryable errors are defined by `is_retryable` (defaults: 5xx / 429 /
    network errors from the provider layer).
  - Non-retryable errors abort the coroutine immediately.
  - Each attempt can record its error via `on_attempt` if the caller wants
    to surface it (e.g. counting retries on a `city_briefs_runs` row).
"""

from __future__ import annotations

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

from app.providers._common import (
    NonRetryableProviderError,
    ProviderHttpError,
    RetryableProviderError,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")

DEFAULT_MAX_ATTEMPTS = 3
DEFAULT_BASE_DELAY = 2.0  # seconds
DEFAULT_MAX_DELAY = 8.0   # seconds


def is_retryable(exc: BaseException) -> bool:
    """Retry only transient, provider-side failures."""
    if isinstance(exc, RetryableProviderError):
        return True
    if isinstance(exc, NonRetryableProviderError):
        return False
    if isinstance(exc, ProviderHttpError):
        return exc.status == 429 or exc.status >= 500
    # httpx transport / network / asyncio timeouts:
    for retryable_type in (asyncio.TimeoutError, ConnectionError):
        if isinstance(exc, retryable_type):
            return True
    return False


async def with_retry(
    awaitable_factory: Callable[[], Awaitable[T]],
    *,
    max_attempts: int = DEFAULT_MAX_ATTEMPTS,
    base_delay: float = DEFAULT_BASE_DELAY,
    max_delay: float = DEFAULT_MAX_DELAY,
    on_attempt: Callable[[int, BaseException | None], Awaitable[None]] | None = None,
    stage_name: str = "",
) -> T:
    """Run `awaitable_factory()` with bounded retries.

    `awaitable_factory` is called fresh on every attempt so async iterators
    can be retried safely.
    """
    attempt = 0
    last_exc: BaseException | None = None
    while attempt < max_attempts:
        attempt += 1
        try:
            result = await awaitable_factory()
            if on_attempt:
                try:
                    await on_attempt(attempt, None)
                except Exception:  # observability must never break the pipeline
                    logger.warning("on_attempt hook raised; suppressing", exc_info=True)
            return result
        except Exception as exc:
            last_exc = exc
            if on_attempt:
                try:
                    await on_attempt(attempt, exc)
                except Exception:
                    logger.warning("on_attempt hook raised; suppressing", exc_info=True)
            if not is_retryable(exc) or attempt >= max_attempts:
                raise
            delay = min(max_delay, base_delay * (2 ** (attempt - 1)))
            delay = delay + random.uniform(0, 0.5)  # jitter
            logger.warning(
                "stage %s attempt %d/%d failed (%s); retrying in %.1fs",
                stage_name or "<unnamed>", attempt, max_attempts,
                type(exc).__name__, delay,
            )
            await asyncio.sleep(delay)
    assert last_exc is not None
    raise last_exc