"""
LLM Client Factory
==================
Unified interface for multiple LLM providers.
Detects the provider from the model name and creates the appropriate client.

Usage:
    from .llm_client import create_llm_client

    client = create_llm_client(model="gpt-4", api_key="sk-...")
    response = client.chat(messages=[...], temperature=0.0)

    client = create_llm_client(model="claude-sonnet-4-20250514", api_key="sk-ant-...")
    response = client.chat(messages=[...], temperature=0.0)

Supported providers:
    - OpenAI:    gpt-*, o1-*, o3-*, chatgpt-*
    - Anthropic: claude-*
    - Google:    gemini-*
    - Ollama:    llama*, mistral*, codellama*, phi-*, qwen*, deepseek* (local)
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
import json
import logging

logger = logging.getLogger(__name__)


# ─────────────────────────────────────────────
# Base Interface
# ─────────────────────────────────────────────

class LLMClient(ABC):
    """Unified interface that all provider clients implement."""

    def __init__(self, model: str, api_key: Optional[str] = None, base_url: Optional[str] = None):
        self.model = model
        self.api_key = api_key
        self.base_url = base_url

    @abstractmethod
    def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.0,
        max_tokens: Optional[int] = None,
        response_format: Optional[Dict[str, Any]] = None,
    ) -> str:
        """
        Send messages and return the text content of the response.

        Args:
            messages: List of {"role": "system"|"user"|"assistant", "content": "..."}
            temperature: Sampling temperature
            max_tokens: Max output tokens
            response_format: Provider-specific response format config

        Returns:
            The text content of the model's response.
        """
        ...

    @property
    @abstractmethod
    def provider_name(self) -> str:
        ...


# ─────────────────────────────────────────────
# OpenAI
# ─────────────────────────────────────────────

class OpenAIClient(LLMClient):
    def __init__(self, model: str, api_key: str, base_url: Optional[str] = None):
        super().__init__(model, api_key, base_url)
        from openai import OpenAI
        kwargs = {"api_key": api_key}
        if base_url:
            kwargs["base_url"] = base_url
        self._client = OpenAI(**kwargs)

    @property
    def provider_name(self) -> str:
        return "openai"

    def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.0,
        max_tokens: Optional[int] = None,
        response_format: Optional[Dict[str, Any]] = None,
    ) -> str:
        from openai import OpenAIError

        kwargs: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
        }

        # Some models (like o1/o3 reasoning) don't support temperature/max_tokens
        if not self.model.startswith(("o1", "o3")):
            kwargs["temperature"] = temperature
            if max_tokens:
                kwargs["max_tokens"] = max_tokens

        if response_format:
            kwargs["response_format"] = response_format

        try:
            response = self._client.chat.completions.create(**kwargs)
        except OpenAIError as e:
            raise RuntimeError(f"OpenAI call failed: {e}") from e

        return response.choices[0].message.content


# ─────────────────────────────────────────────
# Anthropic
# ─────────────────────────────────────────────

class AnthropicClient(LLMClient):
    def __init__(self, model: str, api_key: str, base_url: Optional[str] = None):
        super().__init__(model, api_key, base_url)
        from anthropic import Anthropic
        kwargs = {"api_key": api_key}
        if base_url:
            kwargs["base_url"] = base_url
        self._client = Anthropic(**kwargs)

    @property
    def provider_name(self) -> str:
        return "anthropic"

    def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.0,
        max_tokens: Optional[int] = None,
        response_format: Optional[Dict[str, Any]] = None,
    ) -> str:
        from anthropic import APIError

        # Anthropic separates system from messages
        system_prompt = None
        chat_messages = []
        for msg in messages:
            if msg["role"] == "system":
                system_prompt = msg["content"]
            else:
                chat_messages.append(msg)

        kwargs: Dict[str, Any] = {
            "model": self.model,
            "messages": chat_messages,
            "max_tokens": max_tokens or 4096,
            "temperature": temperature,
        }

        if system_prompt:
            kwargs["system"] = system_prompt

        try:
            response = self._client.messages.create(**kwargs)
        except APIError as e:
            raise RuntimeError(f"Anthropic call failed: {e}") from e

        # Anthropic returns content blocks
        return response.content[0].text


# ─────────────────────────────────────────────
# Google Gemini
# ─────────────────────────────────────────────

class GoogleClient(LLMClient):
    def __init__(self, model: str, api_key: str, base_url: Optional[str] = None):
        super().__init__(model, api_key, base_url)
        from google import genai
        self._genai_client = genai.Client(api_key=api_key)

    @property
    def provider_name(self) -> str:
        return "google"

    def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.0,
        max_tokens: Optional[int] = None,
        response_format: Optional[Dict[str, Any]] = None,
    ) -> str:
        from google.genai import types

        # Convert messages to Gemini format
        system_instruction = None
        contents = []
        for msg in messages:
            if msg["role"] == "system":
                system_instruction = msg["content"]
            elif msg["role"] == "user":
                contents.append(types.Content(
                    role="user",
                    parts=[types.Part(text=msg["content"])]
                ))
            elif msg["role"] == "assistant":
                contents.append(types.Content(
                    role="model",
                    parts=[types.Part(text=msg["content"])]
                ))

        config = types.GenerateContentConfig(
            temperature=temperature,
        )
        if max_tokens:
            config.max_output_tokens = max_tokens
        if system_instruction:
            config.system_instruction = system_instruction

        try:
            response = self._genai_client.models.generate_content(
                model=self.model,
                contents=contents,
                config=config,
            )
        except Exception as e:
            raise RuntimeError(f"Google Gemini call failed: {e}") from e

        return response.text


# ─────────────────────────────────────────────
# Ollama (Local models via OpenAI-compatible API)
# ─────────────────────────────────────────────

class OllamaClient(LLMClient):
    """Uses Ollama's OpenAI-compatible endpoint."""

    DEFAULT_BASE_URL = "http://localhost:11434/v1"

    def __init__(self, model: str, api_key: Optional[str] = None, base_url: Optional[str] = None):
        super().__init__(model, api_key or "ollama", base_url or self.DEFAULT_BASE_URL)
        from openai import OpenAI
        self._client = OpenAI(api_key=self.api_key, base_url=self.base_url)

    @property
    def provider_name(self) -> str:
        return "ollama"

    def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.0,
        max_tokens: Optional[int] = None,
        response_format: Optional[Dict[str, Any]] = None,
    ) -> str:
        kwargs: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
        }
        if max_tokens:
            kwargs["max_tokens"] = max_tokens
        if response_format:
            kwargs["response_format"] = response_format

        try:
            response = self._client.chat.completions.create(**kwargs)
        except Exception as e:
            raise RuntimeError(f"Ollama call failed: {e}") from e

        return response.choices[0].message.content


# ─────────────────────────────────────────────
# Provider Detection & Factory
# ─────────────────────────────────────────────

# Model prefix → provider mapping
_PROVIDER_MAP = {
    # OpenAI
    "gpt-": "openai",
    "o1-": "openai",
    "o3-": "openai",
    "chatgpt-": "openai",
    # Anthropic
    "claude-": "anthropic",
    # Google
    "gemini-": "google",
    # Ollama / local
    "llama": "ollama",
    "mistral": "ollama",
    "codellama": "ollama",
    "phi-": "ollama",
    "qwen": "ollama",
    "deepseek": "ollama",
}

_CLIENT_CLASSES = {
    "openai": OpenAIClient,
    "anthropic": AnthropicClient,
    "google": GoogleClient,
    "ollama": OllamaClient,
}


def detect_provider(model: str) -> str:
    """Detect the provider from the model name."""
    model_lower = model.lower()
    for prefix, provider in _PROVIDER_MAP.items():
        if model_lower.startswith(prefix):
            return provider
    raise ValueError(
        f"Unknown model '{model}'. Cannot detect provider.\n"
        f"Supported prefixes: {list(_PROVIDER_MAP.keys())}"
    )


def create_llm_client(
    model: str,
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
    provider: Optional[str] = None,
) -> LLMClient:
    """
    Create an LLM client for the given model.

    Provider is auto-detected from the model name, or can be forced
    via the `provider` parameter.

    Args:
        model: Model identifier, e.g. "gpt-4", "claude-sonnet-4-20250514", "gemini-pro"
        api_key: API key for the provider
        base_url: Override base URL (useful for proxies or local deployments)
        provider: Force a specific provider instead of auto-detecting

    Returns:
        An LLMClient instance ready to use.
    """
    if provider is None:
        provider = detect_provider(model)

    client_cls = _CLIENT_CLASSES.get(provider)
    if client_cls is None:
        raise ValueError(f"Unknown provider '{provider}'. Supported: {list(_CLIENT_CLASSES.keys())}")

    logger.info(f"Creating {provider} client for model '{model}'")
    return client_cls(model=model, api_key=api_key, base_url=base_url)