from __future__ import annotations

import asyncio
import os
import time
from dataclasses import dataclass, replace
from logging import Logger
from typing import AbstractSet, Callable, FrozenSet, Sequence, TypeVar

import httpx

DEFAULT_RETRYABLE_STATUS_CODES: FrozenSet[int] = frozenset(
    {408, 425, 429, 500, 502, 503, 504}
)
DEFAULT_HTTP_RETRY_ATTEMPTS = 3
DEFAULT_HTTP_RETRY_BASE_SECONDS = 0.5
DEFAULT_HTTP_RETRY_MAX_SECONDS = 5.0
RETRYABLE_TRANSPORT_EXCEPTIONS = (
    httpx.ConnectError,
    httpx.ConnectTimeout,
    httpx.ReadTimeout,
    httpx.WriteTimeout,
    httpx.RemoteProtocolError,
)


@dataclass(frozen=True)
class HttpRetryPolicy:
    attempts: int = DEFAULT_HTTP_RETRY_ATTEMPTS
    base_seconds: float = DEFAULT_HTTP_RETRY_BASE_SECONDS
    max_seconds: float = DEFAULT_HTTP_RETRY_MAX_SECONDS
    retryable_status_codes: AbstractSet[int] = DEFAULT_RETRYABLE_STATUS_CODES

    def delay_seconds(self, attempt: int) -> float:
        return min(self.max_seconds, self.base_seconds * (2 ** (attempt - 1)))

    def with_attempts(self, attempts: int) -> "HttpRetryPolicy":
        return replace(self, attempts=max(1, attempts))


TNumber = TypeVar("TNumber", int, float)


def _read_numeric_env(
    env_names: Sequence[str],
    *,
    parse: Callable[[str], TNumber],
    minimum: TNumber,
    default: TNumber,
    logger: Logger,
) -> TNumber:
    for env_name in env_names:
        value = os.getenv(env_name)
        if not value:
            continue
        try:
            parsed = parse(value)
            return max(minimum, parsed)
        except ValueError:
            logger.warning("Invalid %s=%r, using %s", env_name, value, default)
            return default
    return default


def build_retry_policy(
    *,
    logger: Logger,
    attempts: int | None = None,
    base_seconds: float | None = None,
    max_seconds: float | None = None,
    attempts_env_names: Sequence[str] = ("SNOWGLOBE_HTTP_RETRY_ATTEMPTS",),
    base_seconds_env_names: Sequence[str] = ("SNOWGLOBE_HTTP_RETRY_BASE_SECONDS",),
    max_seconds_env_names: Sequence[str] = ("SNOWGLOBE_HTTP_RETRY_MAX_SECONDS",),
    attempts_default: int = DEFAULT_HTTP_RETRY_ATTEMPTS,
    base_seconds_default: float = DEFAULT_HTTP_RETRY_BASE_SECONDS,
    max_seconds_default: float = DEFAULT_HTTP_RETRY_MAX_SECONDS,
    retryable_status_codes: AbstractSet[int] = DEFAULT_RETRYABLE_STATUS_CODES,
) -> HttpRetryPolicy:
    if attempts is None:
        attempts = _read_numeric_env(
            attempts_env_names,
            parse=int,
            minimum=1,
            default=attempts_default,
            logger=logger,
        )
    else:
        attempts = max(1, attempts)

    if base_seconds is None:
        base_seconds = _read_numeric_env(
            base_seconds_env_names,
            parse=float,
            minimum=0.05,
            default=base_seconds_default,
            logger=logger,
        )
    else:
        base_seconds = max(0.05, base_seconds)

    if max_seconds is None:
        max_seconds = _read_numeric_env(
            max_seconds_env_names,
            parse=float,
            minimum=0.1,
            default=max_seconds_default,
            logger=logger,
        )
    else:
        max_seconds = max(0.1, max_seconds)

    return HttpRetryPolicy(
        attempts=attempts,
        base_seconds=base_seconds,
        max_seconds=max_seconds,
        retryable_status_codes=frozenset(retryable_status_codes),
    )


def retry_policy_from_env(
    *,
    logger: Logger,
    attempts_env_names: Sequence[str] = ("SNOWGLOBE_HTTP_RETRY_ATTEMPTS",),
    base_seconds_env_names: Sequence[str] = ("SNOWGLOBE_HTTP_RETRY_BASE_SECONDS",),
    max_seconds_env_names: Sequence[str] = ("SNOWGLOBE_HTTP_RETRY_MAX_SECONDS",),
    attempts_default: int = DEFAULT_HTTP_RETRY_ATTEMPTS,
    base_seconds_default: float = DEFAULT_HTTP_RETRY_BASE_SECONDS,
    max_seconds_default: float = DEFAULT_HTTP_RETRY_MAX_SECONDS,
    retryable_status_codes: AbstractSet[int] = DEFAULT_RETRYABLE_STATUS_CODES,
) -> HttpRetryPolicy:
    return build_retry_policy(
        logger=logger,
        attempts_env_names=attempts_env_names,
        base_seconds_env_names=base_seconds_env_names,
        max_seconds_env_names=max_seconds_env_names,
        attempts_default=attempts_default,
        base_seconds_default=base_seconds_default,
        max_seconds_default=max_seconds_default,
        retryable_status_codes=retryable_status_codes,
    )


async def async_request_with_retries(
    client: httpx.AsyncClient,
    method: str,
    url: str,
    *,
    policy: HttpRetryPolicy,
    logger: Logger,
    **kwargs,
) -> httpx.Response:
    method_name = method.lower()
    request_fn = getattr(client, method_name, None)
    if request_fn is None:

        async def request_fn(url: str, **kwargs) -> httpx.Response:
            return await client.request(method, url, **kwargs)

    for attempt in range(1, policy.attempts + 1):
        try:
            response = await request_fn(url, **kwargs)
            if (
                attempt < policy.attempts
                and response.status_code in policy.retryable_status_codes
            ):
                delay = policy.delay_seconds(attempt)
                logger.warning(
                    "Retrying %s %s due to status %s (attempt %s/%s, sleep %.2fs)",
                    method,
                    url,
                    response.status_code,
                    attempt,
                    policy.attempts,
                    delay,
                )
                await asyncio.sleep(delay)
                continue
            return response
        except RETRYABLE_TRANSPORT_EXCEPTIONS:
            if attempt >= policy.attempts:
                raise
            delay = policy.delay_seconds(attempt)
            logger.warning(
                "Retrying %s %s after transport error (attempt %s/%s, sleep %.2fs)",
                method,
                url,
                attempt,
                policy.attempts,
                delay,
            )
            await asyncio.sleep(delay)
    raise RuntimeError(f"Exhausted retries for {method} {url}")


def sync_request_with_retries(
    client: httpx.Client,
    method: str,
    url: str,
    *,
    policy: HttpRetryPolicy,
    logger: Logger,
    **kwargs,
) -> httpx.Response:
    method_name = method.lower()
    request_fn = getattr(client, method_name, None)
    if request_fn is None:

        def request_fn(url: str, **kwargs) -> httpx.Response:
            return client.request(method, url, **kwargs)

    for attempt in range(1, policy.attempts + 1):
        try:
            response = request_fn(url, **kwargs)
            if (
                attempt < policy.attempts
                and response.status_code in policy.retryable_status_codes
            ):
                delay = policy.delay_seconds(attempt)
                logger.warning(
                    "Retrying %s %s due to status %s (attempt %s/%s, sleep %.2fs)",
                    method,
                    url,
                    response.status_code,
                    attempt,
                    policy.attempts,
                    delay,
                )
                time.sleep(delay)
                continue
            return response
        except RETRYABLE_TRANSPORT_EXCEPTIONS:
            if attempt >= policy.attempts:
                raise
            delay = policy.delay_seconds(attempt)
            logger.warning(
                "Retrying %s %s after transport error (attempt %s/%s, sleep %.2fs)",
                method,
                url,
                attempt,
                policy.attempts,
                delay,
            )
            time.sleep(delay)
    raise RuntimeError(f"Exhausted retries for {method} {url}")
