"""
DocuMint AI - Resilient HF Inference Client.

Wraps huggingface_hub.InferenceClient with:
  - Per-request timeout (env: HF_TIMEOUT_SECONDS, default 120)
  - Exponential-backoff retry on transient errors (503, 429, timeouts)
  - Configurable max retries (env: HF_MAX_RETRIES, default 2)

Usage:
    from hf_client import create_inference_client, robust_chat_completion

    client = create_inference_client(api_key=token)
    response = robust_chat_completion(client, model=model_id, messages=msgs, max_tokens=2048)
"""

import os
import time
import logging
from typing import Optional

from huggingface_hub import InferenceClient

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

HF_TIMEOUT_SECONDS = int(os.environ.get("HF_TIMEOUT_SECONDS", "120"))
HF_MAX_RETRIES = int(os.environ.get("HF_MAX_RETRIES", "2"))
HF_RETRY_BASE_DELAY = float(os.environ.get("HF_RETRY_BASE_DELAY", "2.0"))

# Transient error indicators (substring match on exception message)
_TRANSIENT_PATTERNS = [
    "503",
    "429",
    "rate limit",
    "service unavailable",
    "temporarily unavailable",
    "overloaded",
    "timeout",
    "timed out",
    "connection",
    "loading",
    "model is currently loading",
]


# ---------------------------------------------------------------------------
# Client factory
# ---------------------------------------------------------------------------

def create_inference_client(
    api_key: str,
    timeout: Optional[int] = None,
    **kwargs,
) -> InferenceClient:
    """Create an InferenceClient with timeout configured.

    Reuse this across calls within a request to share the connection.
    """
    t = timeout or HF_TIMEOUT_SECONDS
    return InferenceClient(
        api_key=api_key,
        timeout=t,
        **kwargs,
    )


# ---------------------------------------------------------------------------
# Retry wrapper
# ---------------------------------------------------------------------------

def _is_transient(exc: Exception) -> bool:
    """Check whether an exception looks transient and worth retrying."""
    msg = str(exc).lower()
    exc_type = type(exc).__name__.lower()
    return any(p in msg or p in exc_type for p in _TRANSIENT_PATTERNS)


def robust_chat_completion(
    client: InferenceClient,
    *,
    max_retries: Optional[int] = None,
    base_delay: Optional[float] = None,
    **kwargs,
):
    """Call client.chat_completion with exponential-backoff retry.

    kwargs are passed directly to chat_completion (model, messages, etc.).
    Only retries on transient errors (503, 429, timeouts).
    Non-transient errors (400, 401, 422) raise immediately.

    Returns the ChatCompletion response on success.
    """
    retries = max_retries if max_retries is not None else HF_MAX_RETRIES
    delay = base_delay or HF_RETRY_BASE_DELAY
    model_id = kwargs.get("model", "unknown")

    last_exc = None
    for attempt in range(1 + retries):
        try:
            return client.chat_completion(**kwargs)
        except Exception as e:
            last_exc = e
            if attempt < retries and _is_transient(e):
                wait = delay * (2 ** attempt)
                logger.warning(
                    f"Transient error on {model_id} (attempt {attempt + 1}/{1 + retries}): "
                    f"{type(e).__name__}: {str(e)[:120]}. Retrying in {wait:.1f}s..."
                )
                time.sleep(wait)
            else:
                # Non-transient or last attempt — let it propagate
                if attempt > 0:
                    logger.warning(
                        f"Giving up on {model_id} after {attempt + 1} attempts: "
                        f"{type(e).__name__}: {str(e)[:120]}"
                    )
                raise

    # Should not reach here, but just in case
    raise last_exc  # type: ignore[misc]