"""
Shared retry utilities with simple backoff primitives.

These are lightweight helpers intended for reuse across provider adapters.
"""

from __future__ import annotations

import asyncio
import random
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Iterable, List, Optional, Set

from g_gremlin.contracts import BatchOutcome


@dataclass
class RetryConfig:
    max_attempts: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    exponential_base: float = 2.0
    jitter: bool = True
    retryable_codes: Set[int] = field(default_factory=lambda: {429, 500, 502, 503})


def _sleep(delay: float) -> None:
    time.sleep(delay)


def _compute_delay(config: RetryConfig, attempt: int) -> float:
    delay = config.base_delay * (config.exponential_base ** (attempt - 1))
    delay = min(delay, config.max_delay)
    if config.jitter:
        delay *= random.uniform(0.5, 1.5)
    return delay


def retry_with_backoff(config: RetryConfig):
    """
    Decorator for retrying sync functions with basic exponential backoff.
    Intended for small surface area (network calls).
    """

    def decorator(func: Callable):
        def wrapper(*args, **kwargs):
            last_exc: Optional[Exception] = None
            for attempt in range(1, config.max_attempts + 1):
                try:
                    return func(*args, **kwargs)
                except Exception as exc:  # pragma: no cover - specific exceptions injected by callers
                    last_exc = exc
                    delay = _compute_delay(config, attempt)
                    if attempt >= config.max_attempts:
                        break
                    _sleep(delay)
            if last_exc:
                raise last_exc
        return wrapper

    return decorator


async def retry_batch_items(
    items: Iterable[Any],
    operation: Callable[[List[Any]], Any],
    config: RetryConfig,
) -> BatchOutcome:
    """
    Retry only failed items for batch operations. The `operation` callable should
    accept a list of items and return a tuple of (succeeded, failed, retryable_failed).
    """
    pending = list(items)
    succeeded: List[Any] = []
    failed: List[Any] = []
    retryable_failed: List[Any] = []

    for attempt in range(1, config.max_attempts + 1):
        if not pending:
            break

        try:
            op_succeeded, op_failed, op_retryable = operation(pending)
        except Exception:  # pragma: no cover - delegate exception surface to caller
            if attempt >= config.max_attempts:
                failed.extend(pending)
                break
            delay = _compute_delay(config, attempt)
            await asyncio.sleep(delay)
            continue

        succeeded.extend(op_succeeded)
        failed.extend(op_failed)
        retryable_failed = list(op_retryable)  # overwrite with most recent retryables

        if not retryable_failed:
            break

        if attempt < config.max_attempts:
            delay = _compute_delay(config, attempt)
            await asyncio.sleep(delay)
            pending = retryable_failed
        else:
            failed.extend(retryable_failed)
            retryable_failed = []

    return BatchOutcome(
        succeeded=succeeded,
        failed=failed,
        retryable_failed=retryable_failed,
    )
