"""LiteLLM wrapper for model-agnostic LLM access.

Supports OpenAI, Anthropic, Google, Mistral, and Ollama via a unified interface.
Implements retry logic with exponential backoff for transient failures.
"""

import logging
import time
from typing import Any

from pydantic import BaseModel

from ingestion.config import IngestionConfig

logger = logging.getLogger(__name__)

# Retry configuration
MAX_RETRIES = 3
BASE_DELAY_SECONDS = 1.0
MAX_DELAY_SECONDS = 30.0
BACKOFF_FACTOR = 2.0


class LLMError(Exception):
    """Raised when an LLM call fails after all retries."""

    def __init__(self, provider: str, model: str, message: str):
        self.provider = provider
        self.model = model
        super().__init__(
            f"LLM call failed [{provider}/{model}]: {message}"
        )


class LLMClient:
    """Model-agnostic LLM client using LiteLLM."""

    def __init__(self, config: IngestionConfig):
        self.model = config.llm_model_string
        self.provider = config.llm_provider
        self.config = config

    def complete(
        self,
        messages: list[dict[str, str]],
        response_format: type[BaseModel] | None = None,
    ) -> str:
        """Send a completion request with retry logic.

        Args:
            messages: List of message dicts with 'role' and 'content' keys.
            response_format: Optional Pydantic model for structured output.

        Returns:
            The assistant's response content as a string.

        Raises:
            LLMError: If all retries are exhausted.
        """
        import litellm

        last_error: Exception | None = None

        for attempt in range(1, MAX_RETRIES + 1):
            try:
                kwargs: dict[str, Any] = {
                    "model": self.model,
                    "messages": messages,
                }

                if response_format is not None:
                    kwargs["response_format"] = response_format

                response = litellm.completion(**kwargs)
                content = response.choices[0].message.content
                return content or ""

            except (
                litellm.RateLimitError,
                litellm.Timeout,
                litellm.ServiceUnavailableError,
                litellm.InternalServerError,
            ) as e:
                last_error = e
                delay = min(
                    BASE_DELAY_SECONDS * (BACKOFF_FACTOR ** (attempt - 1)),
                    MAX_DELAY_SECONDS,
                )
                logger.warning(
                    "LLM call attempt %d/%d failed (%s: %s), retrying in %.1fs",
                    attempt,
                    MAX_RETRIES,
                    type(e).__name__,
                    str(e),
                    delay,
                )
                time.sleep(delay)

            except Exception as e:
                # Non-retryable error
                raise LLMError(
                    provider=self.provider,
                    model=self.model,
                    message=f"{type(e).__name__}: {e}",
                ) from e

        # All retries exhausted
        raise LLMError(
            provider=self.provider,
            model=self.model,
            message=f"All {MAX_RETRIES} retries exhausted. Last error: {last_error}",
        )