"""LLM provider abstraction — unified interface for Ollama, Groq, OpenAI, and Anthropic."""

from __future__ import annotations

import os
from typing import Optional

from langchain_core.language_models import BaseChatModel


# ---------------------------------------------------------------------------
# Ollama Cloud endpoint
# ---------------------------------------------------------------------------

OLLAMA_CLOUD_BASE_URL = "https://ollama.com"

# ---------------------------------------------------------------------------
# Supported models per provider
# ---------------------------------------------------------------------------

OLLAMA_DEFAULT_MODELS = [
    "llama3.2",
    "llama3.1",
    "llama3.3",
    "mistral",
    "gemma2",
    "phi3",
    "qwen2.5",
    "deepseek-r1",
]

GROQ_MODELS = [
    "llama-3.3-70b-versatile",
    "llama-3.1-8b-instant",
    "llama-3.1-70b-versatile",
    "llama3-8b-8192",
    "llama3-70b-8192",
    "mixtral-8x7b-32768",
    "gemma2-9b-it",
]

OPENAI_MODELS = [
    "gpt-4o",
    "gpt-4o-mini",
    "gpt-4-turbo",
    "gpt-3.5-turbo",
]

ANTHROPIC_MODELS = [
    "claude-sonnet-4-20250514",
    "claude-haiku-4-5-20251001",
    "claude-3-5-haiku-20241022",
]


def get_llm(
    provider: str = "ollama",
    model: Optional[str] = None,
    api_key: Optional[str] = None,
    temperature: float = 0.02,
    max_tokens: Optional[int] = None,
    base_url: Optional[str] = None,
) -> BaseChatModel:
    """Create a LangChain chat model from the specified provider.

    Args:
        provider: One of "ollama" (default), "ollama_cloud", "groq", "openai", "anthropic".
            - "ollama": Local Ollama instance (no API key needed, free).
            - "ollama_cloud": Ollama Cloud API (requires OLLAMA_API_KEY).
            - "groq": Groq API (requires GROQ_API_KEY, has free tier).
            - "openai": OpenAI API (requires OPENAI_API_KEY).
            - "anthropic": Anthropic API (requires ANTHROPIC_API_KEY).
        model: Model name. If None, uses the provider's default.
        api_key: API key. Falls back to environment variables.
        temperature: Sampling temperature (default: 0.02 for deterministic output).
        max_tokens: Maximum tokens in response. None = provider default.
        base_url: Custom base URL (overrides defaults for any provider).

    Returns:
        A LangChain BaseChatModel instance.
    """
    provider = provider.lower().strip()

    if provider == "ollama":
        return _get_ollama_local(model, temperature, max_tokens, base_url)
    elif provider == "ollama_cloud":
        return _get_ollama_cloud(model, api_key, temperature, max_tokens, base_url)
    elif provider == "groq":
        return _get_groq(model, api_key, temperature, max_tokens)
    elif provider == "openai":
        return _get_openai(model, api_key, temperature, max_tokens, base_url)
    elif provider == "anthropic":
        return _get_anthropic(model, api_key, temperature, max_tokens)
    else:
        raise ValueError(
            f"Unknown LLM provider '{provider}'. "
            f"Supported: ollama, ollama_cloud, groq, openai, anthropic"
        )


# ---------------------------------------------------------------------------
# Ollama (local) — DEFAULT provider. Free, no API key, no rate limits.
# ---------------------------------------------------------------------------

def _get_ollama_local(
    model: Optional[str],
    temperature: float,
    max_tokens: Optional[int],
    base_url: Optional[str],
) -> BaseChatModel:
    from langchain_ollama import ChatOllama

    model = model or "llama3.2"
    kwargs = dict(
        model=model,
        temperature=temperature,
    )
    if max_tokens:
        kwargs["num_predict"] = max_tokens
    if base_url:
        kwargs["base_url"] = base_url

    return ChatOllama(**kwargs)


# ---------------------------------------------------------------------------
# Ollama Cloud — same models, hosted by Ollama. Requires API key.
# See: https://docs.ollama.com/cloud
# ---------------------------------------------------------------------------

def _get_ollama_cloud(
    model: Optional[str],
    api_key: Optional[str],
    temperature: float,
    max_tokens: Optional[int],
    base_url: Optional[str],
) -> BaseChatModel:
    """Ollama Cloud uses the native Ollama API with Bearer token auth.

    See: https://docs.ollama.com/cloud
    Endpoint: https://ollama.com/api/chat
    """
    from langchain_ollama import ChatOllama
    import httpx

    model = model or "llama3.2"
    api_key = api_key or os.environ.get("OLLAMA_API_KEY")
    if not api_key:
        raise ValueError(
            "Ollama Cloud API key required. "
            "Get one at https://ollama.com/settings/keys and set OLLAMA_API_KEY "
            "or pass api_key=..."
        )

    cloud_url = base_url or OLLAMA_CLOUD_BASE_URL
    auth_headers = {"Authorization": f"Bearer {api_key}"}

    kwargs = dict(
        model=model,
        base_url=cloud_url,
        temperature=temperature,
        client_kwargs={"headers": auth_headers},
    )
    if max_tokens:
        kwargs["num_predict"] = max_tokens

    return ChatOllama(**kwargs)


# ---------------------------------------------------------------------------
# Groq
# ---------------------------------------------------------------------------

def _get_groq(
    model: Optional[str],
    api_key: Optional[str],
    temperature: float,
    max_tokens: Optional[int],
) -> BaseChatModel:
    from langchain_groq import ChatGroq

    model = model or "llama-3.3-70b-versatile"
    api_key = api_key or os.environ.get("GROQ_API_KEY")
    if not api_key:
        raise ValueError("Groq API key required. Set GROQ_API_KEY or pass api_key=...")

    return ChatGroq(
        groq_api_key=api_key,
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
        max_retries=2,
    )


# ---------------------------------------------------------------------------
# OpenAI
# ---------------------------------------------------------------------------

def _get_openai(
    model: Optional[str],
    api_key: Optional[str],
    temperature: float,
    max_tokens: Optional[int],
    base_url: Optional[str],
) -> BaseChatModel:
    from langchain_openai import ChatOpenAI

    model = model or "gpt-4o-mini"
    api_key = api_key or os.environ.get("OPENAI_API_KEY")
    if not api_key:
        raise ValueError("OpenAI API key required. Set OPENAI_API_KEY or pass api_key=...")

    kwargs = dict(
        model=model,
        api_key=api_key,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    if base_url:
        kwargs["base_url"] = base_url

    return ChatOpenAI(**kwargs)


# ---------------------------------------------------------------------------
# Anthropic
# ---------------------------------------------------------------------------

def _get_anthropic(
    model: Optional[str],
    api_key: Optional[str],
    temperature: float,
    max_tokens: Optional[int],
) -> BaseChatModel:
    from langchain_anthropic import ChatAnthropic

    model = model or "claude-sonnet-4-20250514"
    api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        raise ValueError("Anthropic API key required. Set ANTHROPIC_API_KEY or pass api_key=...")

    return ChatAnthropic(
        model=model,
        anthropic_api_key=api_key,
        temperature=temperature,
        max_tokens=max_tokens or 4096,
    )


# ---------------------------------------------------------------------------
# Utility
# ---------------------------------------------------------------------------

def list_providers() -> list[str]:
    """Return available provider names."""
    return ["ollama", "ollama_cloud", "groq", "openai", "anthropic"]


def list_models(provider: str) -> list[str]:
    """Return known model names for a provider."""
    provider = provider.lower().strip()
    return {
        "ollama": OLLAMA_DEFAULT_MODELS,
        "ollama_cloud": OLLAMA_DEFAULT_MODELS,
        "groq": GROQ_MODELS,
        "openai": OPENAI_MODELS,
        "anthropic": ANTHROPIC_MODELS,
    }.get(provider, [])
