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

import json
import os
from typing import Any, AsyncIterator
from loguru import logger

import litellm
from litellm import acompletion

from flowly.providers.base import LLMProvider, LLMResponse, ToolCallRequest
from flowly.providers.key_rotator import KeyRotator, classify_error


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

    Supports OpenRouter, Anthropic, OpenAI, Gemini, and many other providers through
    a unified interface.  When multiple API keys are supplied via `fallback_keys`,
    the provider automatically rotates to the next healthy key on auth errors,
    rate-limit errors, or provider overload.
    """

    def __init__(
        self,
        api_key: str | None = None,
        api_base: str | None = None,
        default_model: str = "anthropic/claude-opus-4-5",
        fallback_keys: list[str] | None = None,
        provider_name: str = "unknown",
    ):
        super().__init__(api_key, api_base)
        self.default_model = default_model
        self.request_timeout_seconds = float(os.getenv("FLOWLY_LLM_TIMEOUT_SECONDS", "90"))
        self.stream_timeout_seconds = float(os.getenv("FLOWLY_LLM_STREAM_TIMEOUT_SECONDS", "120"))

        # 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

        # API keys are passed directly via kwargs in chat() (line 129-130).
        # We intentionally do NOT write keys to os.environ to prevent
        # exposure to child processes and /proc/[pid]/environ.

        if api_base:
            litellm.api_base = api_base

        # Disable LiteLLM logging noise
        litellm.suppress_debug_info = True

        # Key rotation: build list of all keys (primary + fallbacks)
        all_keys: list[str] = []
        if api_key:
            all_keys.append(api_key)
        if fallback_keys:
            all_keys.extend(k for k in fallback_keys if k and k != api_key)

        self._rotator: KeyRotator | None = None
        if len(all_keys) > 1:
            from flowly.audit.logger import get_audit_logger
            audit = get_audit_logger()
            self._rotator = KeyRotator(
                keys=all_keys,
                provider=provider_name,
                on_rotate=lambda prov, reason, fi, ti: audit.log_key_rotation(
                    provider=prov, reason=reason, key_index_from=fi, key_index_to=ti
                ),
            )
            logger.info(
                f"[KeyRotator] {provider_name}: {len(all_keys)} keys configured"
            )
    
    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,
        tool_choice: str = "auto",
    ) -> 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 if not already prefixed
        if self.is_openrouter and not model.startswith("openrouter/"):
            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
        if "gemini" in model.lower() and not model.startswith("gemini/"):
            model = f"gemini/{model}"
        
        kwargs: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "timeout": self.request_timeout_seconds,
        }

        # Set app name for OpenRouter dashboard
        if self.is_openrouter:
            kwargs["extra_headers"] = {
                "X-Title": "Flowly",
                "HTTP-Referer": "https://github.com/pve/flowlyai",
            }
        
        # Pass api_base and api_key directly for custom endpoints
        if self.api_base:
            kwargs["api_base"] = self.api_base
        if self.api_key:
            kwargs["api_key"] = self.api_key
        
        if tools:
            kwargs["tools"] = tools
            # Use provided tool_choice (required/auto/none)
            # "required" forces model to call a tool - prevents hallucination
            # "auto" lets model decide - use after tools have been executed
            kwargs["tool_choice"] = tool_choice
        
        # Attempt with key rotation on retriable errors
        max_attempts = self._rotator.key_count() if self._rotator else 1
        for attempt in range(max_attempts):
            # Inject current key if rotator is active
            if self._rotator:
                kwargs["api_key"] = self._rotator.current_key()
            try:
                response = await acompletion(**kwargs)
                return self._parse_response(response)
            except Exception as e:
                error_msg = str(e)
                # Redact keys from error messages
                for k in ([self.api_key] if self.api_key else []):
                    if k and len(k) > 8:
                        error_msg = error_msg.replace(k, "***")

                rotation_reason = classify_error(error_msg) if self._rotator else None
                if rotation_reason and attempt < max_attempts - 1:
                    logger.warning(
                        f"LLM call failed ({rotation_reason}), rotating key "
                        f"(attempt {attempt + 1}/{max_attempts}): {error_msg[:120]}"
                    )
                    self._rotator.rotate(reason=rotation_reason)
                    continue  # retry with next key

                logger.error(f"LLM call error: {error_msg}")
                return LLMResponse(
                    content=f"Error calling LLM: {error_msg}",
                    finish_reason="error",
                )
    
    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,
        )
    
    async def chat_stream(
        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,
        tool_choice: str = "auto",
    ) -> AsyncIterator[LLMResponse]:
        """
        Stream a chat completion via LiteLLM.

        Yields text-delta LLMResponses as tokens arrive, then a final
        LLMResponse with accumulated tool_calls (if any).
        """
        model = model or self.default_model

        if self.is_openrouter and not model.startswith("openrouter/"):
            model = f"openrouter/{model}"
        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}"
        if self.is_vllm:
            model = f"hosted_vllm/{model}"
        if "gemini" in model.lower() and not model.startswith("gemini/"):
            model = f"gemini/{model}"

        kwargs: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "timeout": self.stream_timeout_seconds,
            "stream": True,
        }
        if self.is_openrouter:
            kwargs["extra_headers"] = {
                "X-Title": "Flowly",
                "HTTP-Referer": "https://github.com/pve/flowlyai",
            }
        if self.api_base:
            kwargs["api_base"] = self.api_base
        if self.api_key:
            kwargs["api_key"] = self.api_key
        if tools:
            kwargs["tools"] = tools
            kwargs["tool_choice"] = tool_choice

        try:
            stream = await acompletion(**kwargs)
        except Exception as e:
            error_msg = str(e)
            if self.api_key and len(self.api_key) > 8:
                error_msg = error_msg.replace(self.api_key, "***")
            logger.error(f"LLM stream error: {error_msg}")
            yield LLMResponse(content=f"Error calling LLM: {error_msg}", finish_reason="error")
            return

        # Accumulate tool call chunks: index → {id, name, arguments_str}
        tool_call_accum: dict[int, dict[str, Any]] = {}
        finish_reason = "stop"

        try:
            async for chunk in stream:
                choice = chunk.choices[0]
                delta = choice.delta
                finish_reason = choice.finish_reason or finish_reason

                # Text delta
                if delta.content:
                    yield LLMResponse(content=delta.content, finish_reason="")

                # Tool call chunks
                if delta.tool_calls:
                    for tc_delta in delta.tool_calls:
                        idx = tc_delta.index
                        if idx not in tool_call_accum:
                            tool_call_accum[idx] = {"id": "", "name": "", "arguments_str": ""}
                        entry = tool_call_accum[idx]
                        if tc_delta.id:
                            entry["id"] = tc_delta.id
                        if tc_delta.function:
                            if tc_delta.function.name:
                                entry["name"] += tc_delta.function.name
                            if tc_delta.function.arguments:
                                entry["arguments_str"] += tc_delta.function.arguments
        except Exception as e:
            error_msg = str(e)
            if self.api_key and len(self.api_key) > 8:
                error_msg = error_msg.replace(self.api_key, "***")
            logger.error(f"LLM stream read error: {error_msg}")

        # Yield final response with tool calls (if any)
        if tool_call_accum:
            tool_calls = []
            for idx in sorted(tool_call_accum):
                entry = tool_call_accum[idx]
                args_str = entry["arguments_str"]
                try:
                    args = json.loads(args_str) if args_str else {}
                except json.JSONDecodeError:
                    args = {"raw": args_str}
                tool_calls.append(ToolCallRequest(
                    id=entry["id"] or f"call_{idx}",
                    name=entry["name"],
                    arguments=args,
                ))
            yield LLMResponse(content=None, tool_calls=tool_calls, finish_reason=finish_reason)
        else:
            # Signal stream end with empty content (so caller knows we're done)
            yield LLMResponse(content=None, finish_reason=finish_reason)

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