"""Retry utilities for LLM API calls.

Provides exponential backoff retry logic for transient errors while
immediately failing on authentication errors.
"""

from __future__ import annotations

import logging
import random
import time
from typing import TYPE_CHECKING, TypeVar

from unfamiliar.llm.exceptions import AuthenticationError
from unfamiliar.llm.providers.factory import get_retry_delay

if TYPE_CHECKING:
    from collections.abc import Callable

logger = logging.getLogger(__name__)

T = TypeVar("T")

# HTTP status codes that indicate transient errors (should retry)
_TRANSIENT_STATUS_CODES = {429, 500, 502, 503, 504}

# HTTP status codes that indicate auth errors (immediate fail)
_AUTH_STATUS_CODES = {401, 403}

# Default retry configuration
_DEFAULT_MAX_ATTEMPTS = 5
_DEFAULT_BASE_DELAY = 1.0


def with_llm_retry(
    fn: Callable[..., T],
    *args: object,
    max_attempts: int = _DEFAULT_MAX_ATTEMPTS,
    base_delay: float | None = None,
    provider_name: str | None = None,
    **kwargs: object,
) -> T:
    """Execute a function with exponential backoff retry on transient LLM errors.

    Args:
        fn: The function to call (typically an LLM chain's .invoke()).
        *args: Positional arguments to pass to fn.
        max_attempts: Maximum number of attempts (1 = no retries).
        base_delay: Base delay in seconds for exponential backoff. If None,
            uses provider-specific delay if provider_name given, else default.
        provider_name: Provider name for provider-specific retry delays.
            "glm" → 30s, "anthropic" → 1s, unknown → 1s.
        **kwargs: Keyword arguments to pass to fn.

    Returns:
        The return value of fn on success.

    Raises:
        AuthenticationError: Immediately on 401/403 errors (no retry).
        Exception: The last exception after max_attempts exhausted.
    """
    # Resolve base_delay: explicit > provider-specific > default
    if base_delay is None:
        if provider_name:
            base_delay = get_retry_delay(provider_name)
        else:
            base_delay = _DEFAULT_BASE_DELAY

    last_error: Exception | None = None

    for attempt in range(1, max_attempts + 1):
        try:
            return fn(*args, **kwargs)
        except Exception as exc:
            last_error = exc

            # Check transient errors first — a 429 should always retry
            # even if the error message happens to mention "api key"
            if _is_transient_error(exc):
                if attempt < max_attempts:
                    delay = _calculate_delay(attempt, base_delay)
                    logger.warning(
                        "Transient error in LLM call (attempt %d/%d). "
                        "Retrying in %.2fs: %s",
                        attempt,
                        max_attempts,
                        delay,
                        exc,
                    )
                    time.sleep(delay)
                    continue
                # Exhausted retries on transient error - log with rate limit info
                rate_limit_info = _extract_rate_limit_info(exc)
                body = _extract_response_body(exc)
                rate_msg = f"LLM call failed after {max_attempts} attempts: {exc}"
                if rate_limit_info:
                    rate_msg += f"\n  Rate limit headers: {rate_limit_info}"
                if body:
                    rate_msg += f"\n  Response body: {body}"
                logger.error(rate_msg)
                break

            # Check for authentication errors - fail immediately
            if _is_auth_error(exc):
                logger.error("Authentication failed for LLM call. Check your API key.")
                raise AuthenticationError(
                    "LLM API authentication failed. Check your API key configuration."
                ) from exc

            # Non-transient, non-auth error - fail immediately
            body = _extract_response_body(exc)
            logger.error(
                "Non-transient error in LLM call (attempt %d/%d): %s%s",
                attempt,
                max_attempts,
                exc,
                f"\n  Response body: {body}" if body else "",
            )
            raise

    # All attempts exhausted
    if last_error is not None:
        raise last_error
    raise RuntimeError("Unexpected state: no error recorded")


def _is_auth_error(exc: Exception) -> bool:
    """Check if an exception indicates an authentication failure."""
    # Check for HTTP status code in various exception types
    status_code = _extract_status_code(exc)
    if status_code is not None:
        return status_code in _AUTH_STATUS_CODES

    # Check for common auth error message patterns
    error_str = str(exc).lower()
    auth_patterns = [
        "unauthorized",
        "forbidden",
        "invalid api key",
        "authentication failed",
    ]
    return any(pattern in error_str for pattern in auth_patterns)


def _is_transient_error(exc: Exception) -> bool:
    """Check if an exception indicates a transient error worth retrying."""
    # Check for HTTP status code
    status_code = _extract_status_code(exc)
    if status_code is not None:
        return status_code in _TRANSIENT_STATUS_CODES

    # Check for connection-related errors
    error_type = type(exc).__name__.lower()
    connection_patterns = [
        "connectionerror",
        "timeout",
        "network",
        "connecttimeout",
        "readtimeout",
    ]
    if any(pattern in error_type for pattern in connection_patterns):
        return True

    # Check for rate limiting or overload in message
    error_str = str(exc).lower()
    transient_patterns = [
        "rate limit",
        "overloaded",
        "too many requests",
        "temporarily unavailable",
        "service unavailable",
        # Schema validation errors - LLM may succeed on retry
        "did not match expected schema",
        "validation error",
    ]
    return any(pattern in error_str for pattern in transient_patterns)


def _extract_response_body(exc: Exception) -> str | None:
    """Extract HTTP response body from an exception or its cause chain.

    httpx.HTTPStatusError includes the full response, but str(exc)
    only shows the status line. LangChain may wrap the original error,
    so we walk __cause__ and __context__ to find it.
    """
    current: BaseException | None = exc
    seen: set[int] = set()
    while current and id(current) not in seen:
        seen.add(id(current))
        response = getattr(current, "response", None)
        if response is not None:
            text = getattr(response, "text", None)
            if text:
                return text[:500]
        # Walk both explicit and implicit cause chains
        current = getattr(current, "__cause__", None) or getattr(
            current, "__context__", None
        )
    return None


def _extract_rate_limit_info(exc: Exception) -> str | None:
    """Extract rate limit info from 429 response headers.

    Returns a human-readable string with Retry-After or rate limit headers,
    or None if not found.
    """
    current: BaseException | None = exc
    seen: set[int] = set()
    while current and id(current) not in seen:
        seen.add(id(current))
        response = getattr(current, "response", None)
        if response is not None:
            headers = getattr(response, "headers", None)
            if headers is not None:
                info_parts = []
                # Check common rate limit headers
                for header in [
                    "Retry-After",
                    "X-RateLimit-Reset",
                    "X-RateLimit-Remaining",
                    "X-RateLimit-Limit",
                    "ratelimit-reset",
                    "ratelimit-remaining",
                ]:
                    value = headers.get(header)
                    if value:
                        info_parts.append(f"{header}: {value}")
                if info_parts:
                    return ", ".join(info_parts)
        current = getattr(current, "__cause__", None) or getattr(
            current, "__context__", None
        )
    return None


def _extract_status_code(exc: Exception) -> int | None:
    """Extract HTTP status code from various exception types."""
    # LangChain/LangGraph exceptions may have status_code attribute
    status_code = getattr(exc, "status_code", None)
    if isinstance(status_code, int):
        return status_code

    # HTTPx/httpcore exceptions
    response = getattr(exc, "response", None)
    if response is not None:
        status_code = getattr(response, "status_code", None)
        if isinstance(status_code, int):
            return status_code

    # Some exceptions store it in .code
    code = getattr(exc, "code", None)
    if code is not None:
        try:
            return int(code)
        except (TypeError, ValueError):
            pass

    return None


def _calculate_delay(attempt: int, base_delay: float) -> float:
    """Calculate delay with exponential backoff and jitter.

    Formula: base_delay * 2^(attempt-1) + random(0, 1)

    Args:
        attempt: Current attempt number (1-indexed).
        base_delay: Base delay in seconds.

    Returns:
        Delay in seconds before the next attempt.
    """
    exponential_delay = base_delay * (2 ** (attempt - 1))
    jitter = random.random()  # 0.0 to 1.0 seconds
    return exponential_delay + jitter
