from __future__ import annotations

import asyncio
import math
import random
import time
from dataclasses import dataclass
from typing import Callable, Set, TypeVar

from ._errors import (
    APIError,
    ConnectError,
    RateLimitError,
    RequestTimeoutError,
)

T = TypeVar("T")

IDEMPOTENT_METHODS: Set[str] = {"GET", "PUT", "DELETE"}
RETRYABLE_STATUS_CODES: Set[int] = {408, 429, 500, 502, 503, 504}


@dataclass
class RetryConfig:
    """Configuration for retry behaviour with exponential backoff."""

    max_retries: int = 2
    initial_delay_s: float = 0.5
    max_delay_s: float = 8.0
    multiplier: float = 2.0


def is_retryable(error: BaseException, method: str) -> bool:
    """Determine whether the error is eligible for automatic retry."""
    if method.upper() not in IDEMPOTENT_METHODS:
        return False

    if isinstance(error, (ConnectError, RequestTimeoutError)):
        return True

    if isinstance(error, APIError):
        return error.status in RETRYABLE_STATUS_CODES

    return False


def compute_delay(attempt: int, config: RetryConfig, error: BaseException) -> float:
    """Compute the delay in seconds before the next retry attempt.

    Uses exponential backoff with +-20% jitter, capped at ``config.max_delay_s``.
    For 429 responses with a ``Retry-After`` header the returned delay is the
    greater of the computed backoff and the server-requested wait time.
    """
    base = config.initial_delay_s * math.pow(config.multiplier, attempt)
    jitter = 1.0 + (random.random() * 0.4 - 0.2)
    delay = min(base * jitter, config.max_delay_s)

    if isinstance(error, RateLimitError) and error.retry_after is not None:
        delay = max(delay, error.retry_after)

    return round(delay, 3)


def with_retry(
    fn: Callable[[], T],
    method: str,
    config: RetryConfig,
) -> T:
    """Execute a synchronous function with automatic retries on transient failures."""
    last_error: BaseException | None = None

    for attempt in range(config.max_retries + 1):
        try:
            return fn()
        except BaseException as error:
            last_error = error
            is_last = attempt == config.max_retries
            if is_last or not is_retryable(error, method):
                raise
            delay = compute_delay(attempt, config, error)
            time.sleep(delay)

    # Unreachable but satisfies the type checker.
    assert last_error is not None
    raise last_error


async def async_with_retry(
    fn: Callable[[], T],
    method: str,
    config: RetryConfig,
) -> T:
    """Execute an async function with automatic retries on transient failures."""
    last_error: BaseException | None = None

    for attempt in range(config.max_retries + 1):
        try:
            result = fn()
            if asyncio.iscoroutine(result):
                return await result  # type: ignore[misc]
            return result  # type: ignore[return-value]
        except BaseException as error:
            last_error = error
            is_last = attempt == config.max_retries
            if is_last or not is_retryable(error, method):
                raise
            delay = compute_delay(attempt, config, error)
            await asyncio.sleep(delay)

    assert last_error is not None
    raise last_error
