"""LiteLLM-based LLM implementation for Fluxibly.

Wraps LiteLLM for 100+ provider support via OpenAI-compatible format.
"""

from __future__ import annotations

from typing import Any

from fluxibly.llm.base import BaseLLM, LLMConfig
from fluxibly.schema.response import (
    ContentItem,
    LLMMetadata,
    LLMOutput,
    LLMResponse,
    TokenUsage,
)
from fluxibly.schema.tools import ToolCall

try:
    import litellm as litellm_lib
except ImportError as e:
    raise ImportError(
        "LiteLLM package is required for LiteLLM. "
        "Install it with: pip install fluxibly[litellm]"
    ) from e


class LiteLLM(BaseLLM):
    """LiteLLM-based LLM — supports 100+ providers via OpenAI-compatible format."""

    def __init__(self, config: LLMConfig) -> None:
        super().__init__(config)
        if config.api_key:
            litellm_lib.api_key = config.api_key
        if config.api_base:
            litellm_lib.api_base = config.api_base

    def prepare(self) -> dict[str, Any]:
        """Prepare OpenAI-compatible request for LiteLLM."""
        prepared: dict[str, Any] = {
            "model": self.config.model,
            "messages": self.messages,
        }

        if self.tools:
            formatted = []
            for tool in self.tools:
                if tool.get("type") == "function":
                    formatted.append(
                        {"type": "function", "function": tool["function"]}
                    )
            if formatted:
                prepared["tools"] = formatted

        if self.config.max_output_tokens:
            prepared["max_tokens"] = self.config.max_output_tokens
        if self.config.temperature is not None:
            prepared["temperature"] = self.config.temperature
        if self.config.top_p is not None:
            prepared["top_p"] = self.config.top_p
        if self.config.stop:
            prepared["stop"] = self.config.stop

        return prepared

    def forward(self, **kwargs: Any) -> LLMResponse:
        """Execute LiteLLM inference."""
        prepared = self.prepare()
        prepared.update(kwargs)

        self.logger.log_input(prepared)

        try:
            raw = litellm_lib.completion(**prepared)
            result = self._parse_output(raw)
            self._apply_pricing(result)
            self.logger.log_output(result)
            return result
        except Exception as e:
            self.logger.log_error(e, prepared)
            raise

    # ── Response Parsing ─────────────────────────────────────────────

    def _parse_output(self, raw: Any) -> LLMResponse:
        """Parse LiteLLM output (OpenAI-compatible format)."""
        if not raw.choices:
            from fluxibly.exceptions import LLMError

            raise LLMError("LiteLLM returned empty response — no choices")
        choice = raw.choices[0]
        message = choice.message

        content_items: list[ContentItem] = []
        tool_calls: list[ToolCall] = []
        text_parts: list[str] = []

        if message.content:
            text_parts.append(message.content)
            content_items.append(
                ContentItem(type="text", text=message.content)
            )

        for tc in message.tool_calls or []:
            tool_call = ToolCall(
                id=tc.id,
                name=tc.function.name,
                arguments=tc.function.arguments,
                type="function",
            )
            tool_calls.append(tool_call)
            content_items.append(
                ContentItem(type="tool_call", tool_call=tool_call)
            )

        return LLMResponse(
            output=LLMOutput(
                output_text="\n".join(text_parts),
                content=content_items,
                tool_calls=tool_calls,
            ),
            metadata=LLMMetadata(
                model=raw.model or self.config.model,
                id=raw.id or "",
                created_at=str(getattr(raw, "created", "")),
                status=(
                    "completed"
                    if choice.finish_reason == "stop"
                    else "incomplete"
                ),
                usage=TokenUsage(
                    input_tokens=getattr(raw.usage, "prompt_tokens", 0),
                    output_tokens=getattr(raw.usage, "completion_tokens", 0),
                    total_tokens=getattr(raw.usage, "total_tokens", 0),
                ),
                stop_reason=self._normalize_stop_reason(choice.finish_reason),
                provider="litellm",
                raw_response=(
                    raw.model_dump() if hasattr(raw, "model_dump") else None
                ),
            ),
        )

    def _normalize_stop_reason(self, finish_reason: str | None) -> str:
        mapping = {
            "stop": "stop",
            "tool_calls": "tool_calls",
            "length": "max_tokens",
            "content_filter": "content_filter",
        }
        return mapping.get(finish_reason or "stop", finish_reason or "stop")
