"""LLM provider adapters.

Each provider implements a query function that accepts a text prompt and
optional images, returning ``(response_text, token_usage)``.
"""

from __future__ import annotations

import base64
import dataclasses
import os
from typing import Any, Dict, List

from consilium.cost import TokenUsage


@dataclasses.dataclass
class ProviderConfig:
    """Configuration for a single LLM provider/model."""

    provider: str  # "openai", "anthropic", "google"
    model: str  # e.g. "gpt-5.2", "claude-sonnet-4-6"
    api_key: str | None = None  # falls back to env var
    temperature: float = 0.7
    max_tokens: int = 4096

    @property
    def display_name(self) -> str:
        return f"{self.provider}/{self.model}"


def _resolve_api_key(provider: str, explicit_key: str | None) -> str:
    """Resolve API key from explicit value or env var."""
    if explicit_key:
        return explicit_key

    env_map = {
        "openai": "OPENAI_API_KEY",
        "anthropic": "ANTHROPIC_API_KEY",
        "google": "GOOGLE_API_KEY",
    }
    env_var = env_map.get(provider)
    if env_var:
        val = os.environ.get(env_var)
        if val:
            return val

    raise ValueError(
        f"No API key found for provider '{provider}'. "
        f"Set {env_map.get(provider, provider.upper() + '_API_KEY')} "
        f"or pass api_key= explicitly."
    )


# ---------------------------------------------------------------------------
# Provider query functions
# ---------------------------------------------------------------------------


def _query_openai(
    config: ProviderConfig,
    prompt: str,
    images: list[bytes] | None = None,
    system: str | None = None,
    json_schema: dict | None = None,
) -> tuple[str, TokenUsage]:
    """Query an OpenAI-compatible model."""
    import openai

    api_key = _resolve_api_key("openai", config.api_key)
    client = openai.OpenAI(api_key=api_key)

    messages: list[dict[str, Any]] = []
    if system:
        messages.append({"role": "system", "content": system})

    content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
    if images:
        for img_bytes in images:
            b64 = base64.b64encode(img_bytes).decode("utf-8")
            content.append(
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{b64}"},
                }
            )
    messages.append({"role": "user", "content": content})

    kwargs: Dict[str, Any] = {
        "model": config.model,
        "messages": messages,
        "temperature": config.temperature,
        "max_tokens": config.max_tokens,
    }
    if json_schema:
        kwargs["response_format"] = {"type": "json_object"}

    resp = client.chat.completions.create(**kwargs)
    text = resp.choices[0].message.content or ""
    usage = TokenUsage(
        model=config.model,
        input_tokens=resp.usage.prompt_tokens if resp.usage else 0,
        output_tokens=resp.usage.completion_tokens if resp.usage else 0,
    )
    return text, usage


def _query_anthropic(
    config: ProviderConfig,
    prompt: str,
    images: list[bytes] | None = None,
    system: str | None = None,
    json_schema: dict | None = None,
) -> tuple[str, TokenUsage]:
    """Query an Anthropic model."""
    import anthropic

    api_key = _resolve_api_key("anthropic", config.api_key)
    client = anthropic.Anthropic(api_key=api_key)

    content: list[dict[str, Any]] = []
    if images:
        for img_bytes in images:
            b64 = base64.b64encode(img_bytes).decode("utf-8")
            content.append(
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/png",
                        "data": b64,
                    },
                }
            )
    content.append({"type": "text", "text": prompt})

    kwargs: Dict[str, Any] = {
        "model": config.model,
        "max_tokens": config.max_tokens,
        "messages": [{"role": "user", "content": content}],
    }
    if system:
        kwargs["system"] = system
    if config.temperature is not None:
        kwargs["temperature"] = config.temperature

    resp = client.messages.create(**kwargs)
    text = resp.content[0].text if resp.content else ""
    usage = TokenUsage(
        model=config.model,
        input_tokens=resp.usage.input_tokens,
        output_tokens=resp.usage.output_tokens,
    )
    return text, usage


def _query_google(
    config: ProviderConfig,
    prompt: str,
    images: list[bytes] | None = None,
    system: str | None = None,
    json_schema: dict | None = None,
) -> tuple[str, TokenUsage]:
    """Query a Google Gemini model."""
    from google import genai
    from google.genai import types

    api_key = _resolve_api_key("google", config.api_key)
    client = genai.Client(api_key=api_key)

    parts: list[Any] = []
    if images:
        for img_bytes in images:
            parts.append(types.Part.from_bytes(data=img_bytes, mime_type="image/png"))
    parts.append(prompt)

    resp = client.models.generate_content(
        model=config.model,
        contents=parts,
        config=types.GenerateContentConfig(
            system_instruction=system,
            temperature=config.temperature,
            max_output_tokens=config.max_tokens,
        ),
    )
    text = resp.text or ""

    input_tokens = 0
    output_tokens = 0
    if hasattr(resp, "usage_metadata") and resp.usage_metadata:
        input_tokens = (
            getattr(resp.usage_metadata, "prompt_token_count", 0) or 0
        )
        output_tokens = (
            getattr(resp.usage_metadata, "candidates_token_count", 0) or 0
        )

    usage = TokenUsage(
        model=config.model,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
    )
    return text, usage


# ---------------------------------------------------------------------------
# Dispatch
# ---------------------------------------------------------------------------

_PROVIDERS = {
    "openai": _query_openai,
    "anthropic": _query_anthropic,
    "google": _query_google,
}


def query_model(
    config: ProviderConfig,
    prompt: str,
    images: list[bytes] | None = None,
    system: str | None = None,
    json_schema: dict | None = None,
) -> tuple[str, TokenUsage]:
    """Dispatch a query to the appropriate provider.

    Args:
        config: Provider/model configuration.
        prompt: The text prompt.
        images: Optional list of image bytes (PNG).
        system: Optional system prompt.
        json_schema: If set, request JSON output (best-effort).

    Returns:
        ``(response_text, TokenUsage)``
    """
    fn = _PROVIDERS.get(config.provider)
    if fn is None:
        raise ValueError(
            f"Unknown provider '{config.provider}'. "
            f"Supported: {sorted(_PROVIDERS.keys())}"
        )
    return fn(
        config, prompt, images=images, system=system, json_schema=json_schema
    )


# ---------------------------------------------------------------------------
# Parse "provider/model" shorthand
# ---------------------------------------------------------------------------

MODEL_ALIASES: Dict[str, tuple[str, str]] = {
    # OpenAI — GPT-5.x series (current flagships, March 2026)
    "gpt-5.2": ("openai", "gpt-5.2"),
    "gpt-5.2-pro": ("openai", "gpt-5.2-pro"),
    "gpt-5": ("openai", "gpt-5"),
    "gpt-5-mini": ("openai", "gpt-5-mini"),
    # OpenAI — GPT-4.x series (non-reasoning)
    "gpt-4.1": ("openai", "gpt-4.1"),
    "gpt-4.1-mini": ("openai", "gpt-4.1-mini"),
    "gpt-4.1-nano": ("openai", "gpt-4.1-nano"),
    # OpenAI — reasoning models
    "o3": ("openai", "o3"),
    "o4-mini": ("openai", "o4-mini"),
    # Anthropic — Claude 4.6 (current flagships, Feb 2026)
    "claude-opus-4-6": ("anthropic", "claude-opus-4-6"),
    "claude-sonnet-4-6": ("anthropic", "claude-sonnet-4-6"),
    # Anthropic — Claude 4.5 (previous gen)
    "claude-haiku-4-5": ("anthropic", "claude-haiku-4-5-20251001"),
    "claude-sonnet-4-5": ("anthropic", "claude-sonnet-4-5-20250514"),
    "claude-opus-4-5": ("anthropic", "claude-opus-4-5"),
    # Google — Gemini 3.x (current flagships, March 2026)
    "gemini-3.1-pro": ("google", "gemini-3.1-pro-preview"),
    "gemini-3-flash": ("google", "gemini-3-flash-preview"),
    # Google — Gemini 2.5 (stable GA)
    "gemini-2.5-pro": ("google", "gemini-2.5-pro"),
    "gemini-2.5-flash": ("google", "gemini-2.5-flash"),
}


def parse_model_string(model_str: str) -> ProviderConfig:
    """Parse a model string like ``"gpt-4.1"`` or ``"openai/gpt-4.1"``
    into a :class:`ProviderConfig`.
    """
    if "/" in model_str:
        provider, model = model_str.split("/", 1)
        return ProviderConfig(provider=provider, model=model)

    if model_str in MODEL_ALIASES:
        provider, model = MODEL_ALIASES[model_str]
        return ProviderConfig(provider=provider, model=model)

    # Infer provider from well-known model name patterns as a fallback.
    _PROVIDER_PATTERNS = [
        ("claude", "anthropic"),
        ("gpt-", "openai"),
        ("o1", "openai"),
        ("o3", "openai"),
        ("o4", "openai"),
        ("gemini", "google"),
    ]
    model_lower = model_str.lower()
    for pattern, provider in _PROVIDER_PATTERNS:
        if pattern in model_lower:
            return ProviderConfig(provider=provider, model=model_str)

    raise ValueError(
        f"Cannot parse model '{model_str}'. Use 'provider/model' format "
        f"or one of: {sorted(MODEL_ALIASES.keys())}"
    )


# Default council composition — latest flagship from each provider
DEFAULT_MODELS: List[str] = [
    "gpt-5.2",
    "claude-sonnet-4-6",
    "gemini-3.1-pro",
]

DEFAULT_CHAIRMAN: str = "claude-sonnet-4-6"


def get_default_models() -> List[str]:
    """Return default model IDs, auto-detecting latest via provider APIs.

    Tries :func:`model_registry.get_default_models` first, falling back
    to the hardcoded :data:`DEFAULT_MODELS` on any error.
    """
    try:
        from consilium.model_registry import (
            get_default_models as _registry_defaults,
        )

        models = _registry_defaults()
        if models:
            return models
    except Exception:
        pass
    return list(DEFAULT_MODELS)
