"""LiteLLM provider implementation for multi-provider support."""

import os
from typing import Any

import litellm
from litellm import acompletion

from debot.providers.base import LLMProvider, LLMResponse, ToolCallRequest


class LiteLLMProvider(LLMProvider):
    """
    LLM provider using LiteLLM for multi-provider support.

    Supports OpenRouter, Anthropic, OpenAI, Gemini, and many other providers through
    a unified interface.
    """

    def __init__(
        self,
        api_key: str | None = None,
        api_base: str | None = None,
        default_model: str = "anthropic/claude-opus-4-5",
        all_api_keys: dict[str, str] | None = None,
    ):
        super().__init__(api_key, api_base)
        self.default_model = default_model

        # Detect OpenRouter by api_key prefix or explicit api_base
        self.is_openrouter = (api_key and api_key.startswith("sk-or-")) or (api_base and "openrouter" in api_base)

        # Track if using custom endpoint (vLLM, etc.)
        self.is_vllm = bool(api_base) and not self.is_openrouter

        # Set ALL configured provider API keys so LiteLLM can route across
        # providers during fallback (e.g. OpenRouter credits exhausted → Anthropic direct).
        _key_env_map = {
            "openrouter": "OPENROUTER_API_KEY",
            "anthropic": "ANTHROPIC_API_KEY",
            "openai": "OPENAI_API_KEY",
            "gemini": "GEMINI_API_KEY",
            "groq": "GROQ_API_KEY",
            "zhipu": "ZHIPUAI_API_KEY",
        }
        if all_api_keys:
            for provider_name, key in all_api_keys.items():
                env_var = _key_env_map.get(provider_name)
                if env_var and key:
                    os.environ.setdefault(env_var, key)

        # Configure primary provider key (overrides setdefault above)
        if api_key:
            if self.is_openrouter:
                os.environ["OPENROUTER_API_KEY"] = api_key
            elif self.is_vllm:
                os.environ["OPENAI_API_KEY"] = api_key
            elif "anthropic" in default_model:
                os.environ["ANTHROPIC_API_KEY"] = api_key
            elif "openai" in default_model or "gpt" in default_model:
                os.environ["OPENAI_API_KEY"] = api_key
            elif "gemini" in default_model.lower():
                os.environ["GEMINI_API_KEY"] = api_key
            elif "zhipu" in default_model or "glm" in default_model or "zai" in default_model:
                os.environ["ZHIPUAI_API_KEY"] = api_key
            elif "groq" in default_model:
                os.environ["GROQ_API_KEY"] = api_key

        # Do NOT set litellm.api_base globally — it would route ALL calls
        # (including cross-provider fallback) through OpenRouter.
        # api_base is passed per-call in chat() kwargs instead.

        # Disable LiteLLM logging noise
        litellm.suppress_debug_info = True

    async def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        model: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> LLMResponse:
        """
        Send a chat completion request via LiteLLM.

        Args:
            messages: List of message dicts with 'role' and 'content'.
            tools: Optional list of tool definitions in OpenAI format.
            model: Model identifier (e.g., 'anthropic/claude-sonnet-4-5').
            max_tokens: Maximum tokens in response.
            temperature: Sampling temperature.

        Returns:
            LLMResponse with content and/or tool calls.
        """
        model = model or self.default_model

        # For OpenRouter, prefix model name — but skip if the model's native
        # provider has a direct API key configured (cross-provider fallback).
        if self.is_openrouter and not model.startswith("openrouter/"):
            # Check if a direct provider key is available for this model
            _direct_available = (
                (model.startswith("anthropic/") and os.environ.get("ANTHROPIC_API_KEY"))
                or (model.startswith("openai/") and os.environ.get("OPENAI_API_KEY"))
                or (model.startswith("groq/") and os.environ.get("GROQ_API_KEY"))
                or (model.startswith("gemini/") and os.environ.get("GEMINI_API_KEY"))
            )
            if not _direct_available:
                model = f"openrouter/{model}"

        # For Zhipu/Z.ai, ensure prefix is present
        # Handle cases like "glm-4.7-flash" -> "zhipu/glm-4.7-flash"
        if ("glm" in model.lower() or "zhipu" in model.lower()) and not (
            model.startswith("zhipu/") or model.startswith("zai/") or model.startswith("openrouter/")
        ):
            model = f"zhipu/{model}"

        # For vLLM, use hosted_vllm/ prefix per LiteLLM docs
        # Convert openai/ prefix to hosted_vllm/ if user specified it
        if self.is_vllm:
            model = f"hosted_vllm/{model}"

        # For Gemini, ensure gemini/ prefix if not already present (skip for OpenRouter)
        if "gemini" in model.lower() and not model.startswith(("gemini/", "openrouter/")):
            model = f"gemini/{model}"

        kwargs: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
        }

        # Pass api_base for custom endpoints — but NOT for direct-provider calls
        # (cross-provider fallback should hit the native API, not OpenRouter).
        _is_direct = not model.startswith("openrouter/") and not model.startswith("hosted_vllm/")
        if self.api_base and not _is_direct:
            kwargs["api_base"] = self.api_base

        if tools:
            kwargs["tools"] = tools
            kwargs["tool_choice"] = "auto"

        try:
            response = await acompletion(**kwargs)
            return self._parse_response(response)
        except Exception as e:
            err_str = str(e).lower()
            # Billing check FIRST — billing errors often contain context-related
            # words like "max_tokens" (e.g. "requires more credits, or fewer
            # max_tokens"), so must be checked before context keywords.
            billing_keywords = (
                "credits",
                "afford",
                "402",
                "billing",
                "payment",
                "quota",
                "budget",
            )
            context_keywords = (
                "context_length",
                "context window",
                "maximum context",
                "token limit",
                "too many tokens",
                "input too long",
                "reduce your prompt",
            )
            if any(kw in err_str for kw in billing_keywords):
                finish_reason = "insufficient_credits"
            elif any(kw in err_str for kw in context_keywords):
                finish_reason = "context_length_exceeded"
            else:
                finish_reason = "error"
            return LLMResponse(
                content=f"Error calling LLM: {str(e)}",
                finish_reason=finish_reason,
            )

    def _parse_response(self, response: Any) -> LLMResponse:
        """Parse LiteLLM response into our standard format."""
        choice = response.choices[0]
        message = choice.message

        tool_calls = []
        if hasattr(message, "tool_calls") and message.tool_calls:
            for tc in message.tool_calls:
                # Parse arguments from JSON string if needed
                args = tc.function.arguments
                if isinstance(args, str):
                    import json

                    try:
                        args = json.loads(args)
                    except json.JSONDecodeError:
                        args = {"raw": args}

                tool_calls.append(
                    ToolCallRequest(
                        id=tc.id,
                        name=tc.function.name,
                        arguments=args,
                    )
                )

        usage = {}
        if hasattr(response, "usage") and response.usage:
            usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            }

        return LLMResponse(
            content=message.content,
            tool_calls=tool_calls,
            finish_reason=choice.finish_reason or "stop",
            usage=usage,
        )

    def get_default_model(self) -> str:
        """Get the default model."""
        return self.default_model
