"""Retry middleware for the ORO SDK.

Provides opt-in retry logic with exponential backoff for transient errors
(429, 502, 503, 504) and network failures.

This is a custom (non-generated) module that survives SDK regeneration
because generate-python.sh backs it up and restores it.
"""

from __future__ import annotations

import asyncio
import math
import random
import time
from typing import Callable, Optional

import httpx
from attrs import define, field


@define
class RetryContext:
    """Context passed to the on_retry callback."""

    attempt: int
    """The current retry attempt number (1-based)."""

    delay_s: float
    """The delay in seconds before this retry."""

    status_code: Optional[int] = None
    """The HTTP status code that triggered the retry, or None for network errors."""

    retry_after_s: Optional[float] = None
    """The Retry-After value from the response header, in seconds."""


DEFAULT_RETRYABLE_STATUSES = frozenset({429, 502, 503, 504})


@define
class RetryConfig:
    """Configuration for retry behavior."""

    max_retries: int = 3
    """Maximum number of retries. Set to 0 to disable."""

    retryable_statuses: frozenset[int] = field(factory=lambda: DEFAULT_RETRYABLE_STATUSES)
    """HTTP status codes that should trigger a retry."""

    base_delay_s: float = 1.0
    """Base delay in seconds for exponential backoff."""

    max_delay_s: float = 30.0
    """Maximum delay in seconds."""

    jitter: bool = True
    """Whether to add random jitter to the delay."""

    on_retry: Optional[Callable[[RetryContext], None]] = None
    """Optional callback invoked before each retry attempt."""


def parse_retry_after(response: httpx.Response) -> Optional[float]:
    """Parse the Retry-After header value into seconds.

    Supports both integer seconds (``120``) and HTTP-date format.

    Returns:
        Retry-After value in seconds, or None if not present/parseable.
    """
    header = response.headers.get("retry-after")
    if header is None:
        return None

    # Try parsing as integer/float seconds
    try:
        seconds = float(header)
        if seconds >= 0:
            return seconds
    except ValueError:
        pass

    # Try parsing as HTTP-date
    try:
        from email.utils import parsedate_to_datetime

        dt = parsedate_to_datetime(header)
        delay = dt.timestamp() - time.time()
        return max(delay, 0.0)
    except (ValueError, TypeError):
        pass

    return None


def compute_delay(
    attempt: int,
    base_delay_s: float,
    max_delay_s: float,
    jitter: bool,
    retry_after_s: Optional[float] = None,
) -> float:
    """Compute the delay before the next retry using exponential backoff.

    Formula: ``min(base * 2^attempt, max)``, optionally jittered, with
    Retry-After taking precedence when present.
    """
    if retry_after_s is not None:
        return min(retry_after_s, max_delay_s)

    exponential = base_delay_s * math.pow(2, attempt)
    capped = min(exponential, max_delay_s)

    if jitter:
        return random.random() * capped

    return capped


def _retry_delay_for_response(
    cfg: RetryConfig, attempt: int, response: httpx.Response
) -> Optional[float]:
    """Return delay if the response should be retried, None otherwise."""
    if attempt >= cfg.max_retries or response.status_code not in cfg.retryable_statuses:
        return None
    retry_after = parse_retry_after(response)
    delay = compute_delay(attempt, cfg.base_delay_s, cfg.max_delay_s, cfg.jitter, retry_after)
    if cfg.on_retry is not None:
        cfg.on_retry(RetryContext(
            attempt=attempt + 1,
            delay_s=delay,
            status_code=response.status_code,
            retry_after_s=retry_after,
        ))
    return delay


def _retry_delay_for_error(cfg: RetryConfig, attempt: int) -> Optional[float]:
    """Return delay if a network error should be retried, None otherwise."""
    if attempt >= cfg.max_retries:
        return None
    delay = compute_delay(attempt, cfg.base_delay_s, cfg.max_delay_s, cfg.jitter)
    if cfg.on_retry is not None:
        cfg.on_retry(RetryContext(attempt=attempt + 1, delay_s=delay))
    return delay


class RetryTransport(httpx.BaseTransport):
    """Sync httpx transport that retries transient failures with exponential backoff.

    The transport stacking order is: RetryTransport -> SigningTransport -> HTTPTransport.
    Each retry re-invokes the inner transport, so auth headers are re-signed fresh.
    """

    def __init__(
        self,
        config: RetryConfig,
        wrapped: httpx.BaseTransport | None = None,
    ):
        self._config = config
        self._wrapped = wrapped or httpx.HTTPTransport()

    def handle_request(self, request: httpx.Request) -> httpx.Response:
        last_error: BaseException | None = None
        last_response: httpx.Response | None = None

        for attempt in range(self._config.max_retries + 1):
            try:
                response = self._wrapped.handle_request(request)
                delay = _retry_delay_for_response(self._config, attempt, response)
                if delay is not None:
                    last_response = response
                    time.sleep(delay)
                    continue
                return response
            except (httpx.ConnectError, httpx.ConnectTimeout) as exc:
                last_error = exc
                delay = _retry_delay_for_error(self._config, attempt)
                if delay is not None:
                    time.sleep(delay)
                    continue

        if last_response is not None:
            return last_response
        raise last_error  # type: ignore[misc]


class AsyncRetryTransport(httpx.AsyncBaseTransport):
    """Async httpx transport that retries transient failures with exponential backoff.

    Same behavior as ``RetryTransport`` but uses ``asyncio.sleep()`` and
    ``handle_async_request()``.
    """

    def __init__(
        self,
        config: RetryConfig,
        wrapped: httpx.AsyncBaseTransport | None = None,
    ):
        self._config = config
        self._wrapped = wrapped or httpx.AsyncHTTPTransport()

    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
        last_error: BaseException | None = None
        last_response: httpx.Response | None = None

        for attempt in range(self._config.max_retries + 1):
            try:
                response = await self._wrapped.handle_async_request(request)
                delay = _retry_delay_for_response(self._config, attempt, response)
                if delay is not None:
                    last_response = response
                    await asyncio.sleep(delay)
                    continue
                return response
            except (httpx.ConnectError, httpx.ConnectTimeout) as exc:
                last_error = exc
                delay = _retry_delay_for_error(self._config, attempt)
                if delay is not None:
                    await asyncio.sleep(delay)
                    continue

        if last_response is not None:
            return last_response
        raise last_error  # type: ignore[misc]
