"""
LiteLLM provider implementation.

This implementation uses LiteLLM to provide unified access to 100+ LLM providers
through a single interface. Supports OpenAI, Anthropic, local models (Ollama),
and many more without code changes.

LiteLLM handles:
- API translation across providers
- Retry logic and fallbacks
- Cost tracking and token usage
- Consistent error handling
see docs: https://docs.litellm.ai/docs/
"""

import logging
from typing import Any, Dict, Iterator, List, Optional

from mini_videorag.core.llm import BaseLLM, LLMMessage, LLMResponse

try:
    from litellm import completion, completion_cost, embedding, token_counter
except ImportError as exc:
    raise ImportError("LiteLLM is not installed... Install with: pip install litellm") from exc


class LiteLLMLLM(BaseLLM):
    """
    LLM implementation using LiteLLM for multi-provider support.

    This is the recommended LLM implementation as it provides:
    - 100+ LLM providers through one interface
    - Automatic retry and fallback logic
    - Cost tracking and token usage
    - Support for local models (Ollama, vLLM)
    - Production-ready error handling

    Examples:
        # OpenAI GPT-4
        llm = LiteLLMLLM(model="gpt-4")

        # Anthropic Claude
        llm = LiteLLMLLM(model="claude-3-opus-20240229")

        # Local Ollama
        llm = LiteLLMLLM(model="ollama/llama3")

        # With fallback strategy
        llm = LiteLLMLLM(
            model="gpt-4",
            fallbacks=["gpt-3.5-turbo", "claude-3-sonnet"]
        )

        # With cost limits
        llm = LiteLLMLLM(
            model="gpt-4",
            max_budget=10.0  # dollars
        )

    Environment Variables:
        Model-specific API keys are read automatically:
        - OPENAI_API_KEY: For OpenAI models
        - ANTHROPIC_API_KEY: For Anthropic models
        - COHERE_API_KEY: For Cohere models
        - etc.

    Attributes:
        model: Model identifier (e.g., "gpt-4", "claude-3-opus")
        fallbacks: List of fallback models if primary fails
        default_params: Default parameters for all completions
        logger: Logger instance
    """

    def __init__(
        self,
        model: str,
        fallbacks: Optional[List[str]] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs: Any,
    ):
        """
        Initialize LiteLLM provider.

        Args:
            model: Model identifier (e.g., "gpt-4", "claude-3-opus-20240229")
                   See https://docs.litellm.ai/docs/providers for full list
            fallbacks: Optional list of fallback models to try if primary fails
            temperature: Default temperature (0.0-1.0). Can be overridden per call
            max_tokens: Default max tokens. Can be overridden per call
            **kwargs: Additional default parameters passed to LiteLLM
                     (top_p, frequency_penalty, presence_penalty, etc.)
        """
        self.model = model
        self.fallbacks = fallbacks or []
        self.logger = logging.getLogger(__name__)

        self.default_params: Dict[str, Any] = {}
        if temperature is not None:
            self.default_params["temperature"] = temperature
        if max_tokens is not None:
            self.default_params["max_tokens"] = max_tokens
        self.default_params.update(kwargs)

        self.logger.info(
            f"Initialized LiteLLM: model={model}, "
            f"fallbacks={fallbacks}, params={self.default_params}"
        )

    @property
    def model_name(self) -> str:
        """
        Return the primary model identifier.

        Returns:
            Model name string
        """
        return self.model

    def _convert_messages(self, messages: List[LLMMessage]) -> List[Dict[str, str]]:
        """
        Convert our LLMMessage format to LiteLLM format.

        Args:
            messages: List of LLMMessage objects

        Returns:
            List of dicts in LiteLLM/OpenAI format
        """
        return [{"role": msg.role, "content": msg.content} for msg in messages]

    def _merge_params(self, **kwargs: Any) -> Dict[str, Any]:
        """
        Merge default parameters with call-specific parameters.

        Call-specific parameters override defaults.

        Args:
            **kwargs: Call-specific parameters

        Returns:
            Merged parameter dictionary
        """
        params = self.default_params.copy()
        params.update(kwargs)
        return params

    def generate(
        self,
        messages: List[LLMMessage],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        **kwargs: Any,
    ) -> LLMResponse:
        """
        Generate a response using LiteLLM.

        This method calls LiteLLM's completion endpoint and converts
        the response to our standard LLMResponse format.

        Args:
            messages: Conversation history
            temperature: Sampling temperature (0.0 = deterministic, 1.0 = creative)
            max_tokens: Maximum tokens to generate
            **kwargs: Additional provider-specific parameters

        Returns:
            LLMResponse with generated content and metadata

        Raises:
            Exception: If generation fails (wrapped LiteLLM exceptions)
        """
        self.validate_messages(messages)
        litellm_messages = self._convert_messages(messages)
        params = self._merge_params(temperature=temperature, max_tokens=max_tokens, **kwargs)

        try:
            self.logger.debug(
                f"Calling LiteLLM completion: model={self.model}, "
                f"messages={len(messages)}, params={params}"
            )

            response = completion(
                model=self.model,
                messages=litellm_messages,
                fallbacks=self.fallbacks if self.fallbacks else None,
                **params,
            )

            response_cost = None
            try:
                response_cost = completion_cost(completion_response=response)
            except Exception:  # pylint: disable=broad-exception-caught
                response_cost = None

            llm_response = LLMResponse(
                content=response.choices[0].message.content,
                model=response.model,
                usage={
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens,
                },
                finish_reason=response.choices[0].finish_reason,
                metadata={
                    "response_id": response.id,
                    "created": response.created,
                    "response_cost": response_cost,
                },
            )

            self.logger.info(
                f"LiteLLM completion successful: model={response.model}, "
                f"tokens={response.usage.total_tokens}"
            )

            return llm_response

        except Exception as e:
            self.logger.error(f"LiteLLM completion failed: {e}", exc_info=True)
            raise

    def generate_stream(
        self,
        messages: List[LLMMessage],
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        **kwargs: Any,
    ) -> Iterator[str]:
        """
        Generate a streaming response using LiteLLM.

        Yields content chunks as they are generated, enabling real-time display.

        Args:
            messages: Conversation history
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            **kwargs: Additional provider-specific parameters

        Yields:
            String chunks of generated content

        Raises:
            Exception: If generation fails
        """
        self.validate_messages(messages)
        litellm_messages = self._convert_messages(messages)
        params = self._merge_params(
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True,  # Enable streaming
            **kwargs,
        )

        try:
            self.logger.debug(f"Starting LiteLLM streaming: model={self.model}")

            response = completion(
                model=self.model,
                messages=litellm_messages,
                fallbacks=self.fallbacks if self.fallbacks else None,
                **params,
            )

            for chunk in response:
                if hasattr(chunk.choices[0], "delta"):
                    delta = chunk.choices[0].delta
                    if hasattr(delta, "content") and delta.content:
                        yield delta.content

        except Exception as e:
            self.logger.error(f"LiteLLM streaming failed: {e}", exc_info=True)
            raise

    @property
    def supports_streaming(self) -> bool:
        """
        LiteLLM supports streaming for most providers.

        Returns:
            True
        """
        return True

    @property
    def supports_embeddings(self) -> bool:
        """
        LiteLLM supports embeddings through the embedding() function.

        Returns:
            True
        """
        return True

    def embed(self, text: str) -> List[float]:
        """
        Generate an embedding vector for text using LiteLLM.

        Uses the appropriate embedding model for the provider.
        For example:
        - OpenAI: text-embedding-ada-002
        - Cohere: embed-english-v3.0
        - etc.

        Args:
            text: Text to embed

        Returns:
            Embedding vector as list of floats

        Raises:
            ImportError: If LiteLLM is not installed
            Exception: If embedding generation fails
        """
        try:
            embedding_model = self._get_embedding_model()

            self.logger.debug(f"Generating embedding with {embedding_model}")

            response = embedding(model=embedding_model, input=[text])

            return response.data[0]["embedding"]

        except Exception as e:
            self.logger.error(f"Embedding generation failed: {e}", exc_info=True)
            raise

    def _get_embedding_model(self) -> str:
        """
        Get the appropriate embedding model for this provider.

        Maps chat models to their embedding counterparts.

        Returns:
            Embedding model identifier
        """
        if "gpt" in self.model.lower() or "openai" in self.model.lower():
            return "text-embedding-ada-002"
        if "claude" in self.model.lower():
            # Anthropic doesn't have embeddings, use OpenAI as fallback
            return "text-embedding-ada-002"
        if "cohere" in self.model.lower():
            return "embed-english-v3.0"
        # Default to OpenAI
        return "text-embedding-ada-002"

    @property
    def max_context_length(self) -> int:
        """
        Get maximum context length for this model.

        This is a best-effort mapping. For precise values,
        check LiteLLM's model documentation.

        Returns:
            Maximum context length in tokens
        """
        context_map = {
            "gpt-4": 8192,
            "gpt-4-32k": 32768,
            "gpt-4-turbo": 128000,
            "gpt-3.5-turbo": 16385,
            "claude-3-opus": 200000,
            "claude-3-sonnet": 200000,
            "claude-3-haiku": 200000,
            "claude-2": 100000,
            "llama-3-70b": 8192,
            "mixtral": 32768,
        }

        model_lower = self.model.lower()
        for key, length in context_map.items():
            if key in model_lower:
                return length

        return 4096

    def count_tokens(self, text: str) -> int:
        """
        Estimate token count using LiteLLM's token counter.

        This provides accurate counts for most providers.

        Args:
            text: Text to count tokens for

        Returns:
            Token count
        """
        try:
            return token_counter(model=self.model, text=text)
        except ImportError:
            return super().count_tokens(text)
        except Exception:  # pylint: disable=broad-exception-caught
            return super().count_tokens(text)
