"""
Lightweight multi-provider LLM client for Aigie's internal judge.

Hybrid approach: tries provider SDK first (best latency via connection
pooling + HTTP/2), falls back to raw httpx (zero extra dependencies).

Supported providers (all expose OpenAI-compatible APIs):
    gemini/*      → Google Generative AI
    groq/*        → Groq Cloud
    together/*    → Together AI
    fireworks/*   → Fireworks AI
    mistral/*     → Mistral AI
    openai/*      → OpenAI (also the default)
    openrouter/*  → OpenRouter
    anthropic/*   → Anthropic (Messages API, different format)
"""

import asyncio
import logging
from typing import Any, Dict, Optional

import httpx

logger = logging.getLogger("aigie.judge.llm_client")

# Provider → base URL for OpenAI-compatible chat/completions
_PROVIDER_URLS: Dict[str, str] = {
    "gemini": "https://generativelanguage.googleapis.com/v1beta/openai",
    "groq": "https://api.groq.com/openai/v1",
    "together": "https://api.together.xyz/v1",
    "fireworks": "https://api.fireworks.ai/inference/v1",
    "mistral": "https://api.mistral.ai/v1",
    "openai": "https://api.openai.com/v1",
    "openrouter": "https://openrouter.ai/api/v1",
}

_ANTHROPIC_URL = "https://api.anthropic.com/v1/messages"
_ANTHROPIC_VERSION = "2023-06-01"


def parse_model_string(model: str) -> tuple[str, str]:
    """Parse 'provider/model-name' → (provider, model_name).

    If no slash, assumes 'openai' provider.
    """
    if "/" in model:
        provider, model_name = model.split("/", 1)
        return provider.lower(), model_name
    return "openai", model


class JudgeLLMClient:
    """Multi-provider LLM client for the judge.

    Hybrid strategy:
    1. If the provider's SDK is installed, use it (best latency — connection
       pooling, HTTP/2, retries built in).
    2. Otherwise, fall back to httpx with a persistent AsyncClient (reuses
       TCP+TLS connections, ~150ms saving per call after the first).
    """

    def __init__(self, api_key: str, model: str, timeout: float = 30.0):
        self.api_key = api_key
        self.model = model
        self.timeout = timeout

        self._provider, self._model_name = parse_model_string(model)
        self._sdk_client: Optional[Any] = None
        self._http_client: Optional[httpx.AsyncClient] = None
        self._use_sdk = False

        # Try to create SDK client (best latency)
        self._sdk_client = self._try_create_sdk_client()
        if self._sdk_client:
            self._use_sdk = True
            logger.info(f"Judge LLM client: {self._provider} SDK (model={self._model_name})")
        else:
            logger.info(
                f"Judge LLM client: httpx direct "
                f"(provider={self._provider}, model={self._model_name})"
            )

    def _try_create_sdk_client(self) -> Optional[Any]:
        """Try to create a provider SDK client. Returns None if SDK not installed."""
        provider = self._provider

        if provider in ("openai", "groq", "together", "fireworks", "mistral", "openrouter"):
            # All these use the openai SDK with a custom base_url
            try:
                from openai import OpenAI

                base_url = _PROVIDER_URLS.get(provider)
                if provider == "openai":
                    return OpenAI(api_key=self.api_key)
                return OpenAI(api_key=self.api_key, base_url=base_url + "/")
            except ImportError:
                return None

        if provider == "gemini":
            try:
                from openai import OpenAI

                return OpenAI(
                    api_key=self.api_key,
                    base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
                )
            except ImportError:
                return None

        if provider == "anthropic":
            try:
                import anthropic

                return anthropic.Anthropic(api_key=self.api_key)
            except ImportError:
                return None

        return None

    def _get_http_client(self) -> httpx.AsyncClient:
        """Get or create a persistent httpx client for connection reuse."""
        if self._http_client is None or self._http_client.is_closed:
            self._http_client = httpx.AsyncClient(
                timeout=self.timeout,
                limits=httpx.Limits(
                    max_connections=5,
                    max_keepalive_connections=2,
                    keepalive_expiry=120,  # Keep connections alive for 2 min
                ),
            )
        return self._http_client

    async def complete(self, prompt: str, temperature: float = 0.1, max_tokens: int = 1000) -> str:
        """Send a chat completion request and return the text response."""
        if self._use_sdk and self._sdk_client:
            return await self._call_via_sdk(prompt, temperature, max_tokens)
        return await self._call_via_httpx(prompt, temperature, max_tokens)

    async def _call_via_sdk(self, prompt: str, temperature: float, max_tokens: int) -> str:
        """Call LLM via provider SDK (sync SDK in executor to not block)."""
        loop = asyncio.get_event_loop()

        if self._provider == "anthropic":
            response = await loop.run_in_executor(
                None,
                lambda: self._sdk_client.messages.create(
                    model=self._model_name,
                    max_tokens=max_tokens,
                    messages=[{"role": "user", "content": prompt}],
                    temperature=temperature,
                ),
            )
            return response.content[0].text
        else:
            # OpenAI-compatible SDK
            response = await loop.run_in_executor(
                None,
                lambda: self._sdk_client.chat.completions.create(
                    model=self._model_name,
                    messages=[{"role": "user", "content": prompt}],
                    temperature=temperature,
                    max_tokens=max_tokens,
                ),
            )
            return response.choices[0].message.content

    async def _call_via_httpx(self, prompt: str, temperature: float, max_tokens: int) -> str:
        """Call LLM via raw httpx (fallback when SDK not installed)."""
        if self._provider == "anthropic":
            return await self._httpx_anthropic(prompt, temperature, max_tokens)
        return await self._httpx_openai_compat(prompt, temperature, max_tokens)

    async def _httpx_openai_compat(self, prompt: str, temperature: float, max_tokens: int) -> str:
        """Call any OpenAI-compatible endpoint via httpx."""
        base_url = _PROVIDER_URLS.get(self._provider)
        if not base_url:
            raise ValueError(
                f"Unknown provider '{self._provider}'. "
                f"Supported: {', '.join(sorted(_PROVIDER_URLS.keys()))}"
            )

        client = self._get_http_client()
        response = await client.post(
            f"{base_url}/chat/completions",
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": self._model_name,
                "messages": [{"role": "user", "content": prompt}],
                "temperature": temperature,
                "max_tokens": max_tokens,
            },
        )
        response.raise_for_status()
        data = response.json()

        try:
            return data["choices"][0]["message"]["content"]
        except (KeyError, IndexError) as e:
            raise ValueError(f"Could not parse {self._provider} response: {data}") from e

    async def _httpx_anthropic(self, prompt: str, temperature: float, max_tokens: int) -> str:
        """Call Anthropic Messages API via httpx."""
        client = self._get_http_client()
        response = await client.post(
            _ANTHROPIC_URL,
            headers={
                "x-api-key": self.api_key,
                "anthropic-version": _ANTHROPIC_VERSION,
                "Content-Type": "application/json",
            },
            json={
                "model": self._model_name,
                "max_tokens": max_tokens,
                "messages": [{"role": "user", "content": prompt}],
                "temperature": temperature,
            },
        )
        response.raise_for_status()
        data = response.json()

        try:
            return data["content"][0]["text"]
        except (KeyError, IndexError) as e:
            raise ValueError(f"Could not parse Anthropic response: {data}") from e

    async def close(self) -> None:
        """Close the persistent HTTP client."""
        if self._http_client and not self._http_client.is_closed:
            await self._http_client.aclose()
