"""Shared LLM utilities with multi-provider support.

Provider priority (auto-detected, zero config):
  1. Ollama (local)  — if Ollama is running and has a model
  2. Gemini (cloud)  — if GOOGLE_API_KEY is set
  3. OpenAI-compat   — if OPENAI_API_KEY is set
  4. None            — skip AI, show helpful setup message

Override via:
  - ORBIT_AI_PROVIDER env var (ollama, gemini, openai)
  - ORBIT_AI_MODEL env var (gemma4, gemini-2.5-flash, etc.)
  - orbit auth set ai-provider ollama
  - orbit auth set ai-model gemma4
"""

from __future__ import annotations

import base64
import hashlib
import json
import logging
import os
import urllib.request
from dataclasses import dataclass
from pathlib import Path

logger = logging.getLogger(__name__)

CACHE_DIR = Path.home() / ".orbit" / "cache"

_GEMINI_URL = (
    "https://generativelanguage.googleapis.com/v1beta/"
    "models/{model}:generateContent?key={key}"
)

OLLAMA_BASE_URL = "http://localhost:11434"

# Default models per provider
_DEFAULT_MODELS = {
    "ollama": "gemma4",
    "gemini": "gemini-2.5-flash-lite",
    "openai": "gpt-4o-mini",
}

# Models known to support vision
_VISION_MODELS = {
    "gemma4", "gemma4:e4b", "gemma4:12b", "gemma4:27b",
    "llava", "llava:13b", "llava:34b",
    "llama3.2-vision", "llama3.2-vision:11b", "llama3.2-vision:90b",
    "moondream", "bakllava",
}

# Preferred model order for Ollama auto-selection
_OLLAMA_MODEL_PREFERENCE = [
    "gemma4", "gemma4:e4b", "gemma4:12b", "gemma4:27b",
    "llama3.2-vision", "qwen2.5", "qwen3", "llama3",
    "mistral", "phi3",
]


# ---------------------------------------------------------------------------
# Provider detection
# ---------------------------------------------------------------------------


@dataclass
class ProviderInfo:
    """Detected AI provider."""
    provider: str           # "ollama", "gemini", "openai", or "none"
    model: str              # model name to use
    base_url: str | None    # for openai-compatible providers
    api_key: str | None     # API key (or "ollama" for local)
    supports_vision: bool   # whether the model handles images
    display_name: str       # human-readable description


_cached_provider: ProviderInfo | None = None
_cached_provider_force: str | None = None  # tracks what override was used


def _load_credentials() -> dict:
    """Load stored credentials from ~/.config/orbit/credentials.json."""
    creds_path = Path.home() / ".config" / "orbit" / "credentials.json"
    if creds_path.exists():
        try:
            return json.loads(creds_path.read_text())
        except (json.JSONDecodeError, OSError):
            pass
    return {}


def _get_config_value(key: str) -> str | None:
    """Get a config value from env var, credentials, or orbit config."""
    # 1. Environment variable (highest priority)
    env_map = {
        "ai_provider": "ORBIT_AI_PROVIDER",
        "ai_model": "ORBIT_AI_MODEL",
    }
    env_var = env_map.get(key)
    if env_var:
        val = os.environ.get(env_var)
        if val:
            return val

    # 2. Credentials file
    creds = _load_credentials()
    if key in creds:
        return creds[key]

    # 3. Orbit config
    try:
        from orbit.config import get_config
        cfg = get_config()
        val = getattr(cfg, key, None)
        if val:
            return val
    except Exception:
        pass

    return None


def load_api_key() -> str | None:
    """Load GOOGLE_API_KEY from environment, orbit auth store, or backend/.env."""
    key = os.environ.get("GOOGLE_API_KEY")
    if key:
        return key

    creds = _load_credentials()
    stored_key = creds.get("google_api_key")
    if stored_key:
        return stored_key

    # Search common .env locations
    candidates = [
        Path("backend/.env"),
        Path(__file__).parent.parent.parent / "backend" / ".env",
    ]
    for env_path in candidates:
        if env_path.exists():
            try:
                for line in env_path.read_text().splitlines():
                    line = line.strip()
                    if line.startswith("GOOGLE_API_KEY="):
                        val = line.split("=", 1)[1].strip().strip('"').strip("'")
                        if val:
                            return val
            except OSError:
                continue
    return None


def check_ollama() -> tuple[bool, list[str]]:
    """Check if Ollama is running and return (running, model_names).

    Uses a 2-second timeout. Returns (False, []) on any failure.
    """
    try:
        req = urllib.request.Request(f"{OLLAMA_BASE_URL}/api/tags")
        with urllib.request.urlopen(req, timeout=2) as resp:
            data = json.loads(resp.read())
        models = data.get("models", [])
        names = [m.get("name", "").split(":")[0] or m.get("name", "") for m in models]
        # Also keep full names with tags
        full_names = [m.get("name", "") for m in models]
        # Deduplicate while preserving order
        seen = set()
        unique = []
        for n in names + full_names:
            if n and n not in seen:
                seen.add(n)
                unique.append(n)
        return True, unique
    except Exception:
        return False, []


def _pick_ollama_model(available: list[str], need_vision: bool = False) -> str | None:
    """Pick the best Ollama model from available list."""
    if not available:
        return None

    if need_vision:
        # Prefer vision-capable models
        for pref in _OLLAMA_MODEL_PREFERENCE:
            for avail in available:
                base = avail.split(":")[0]
                if base == pref and (base in _VISION_MODELS or avail in _VISION_MODELS):
                    return avail
        # Check all available for vision support
        for avail in available:
            base = avail.split(":")[0]
            if base in _VISION_MODELS or avail in _VISION_MODELS:
                return avail
        return None  # No vision model found

    # Text: prefer by priority order
    for pref in _OLLAMA_MODEL_PREFERENCE:
        for avail in available:
            if avail.split(":")[0] == pref or avail == pref:
                return avail
    # Fall back to first available
    return available[0] if available else None


def detect_provider(need_vision: bool = False) -> ProviderInfo:
    """Auto-detect the best available AI provider.

    Results are cached for the process lifetime unless override changes.
    """
    global _cached_provider, _cached_provider_force

    force_provider = _get_config_value("ai_provider")
    force_model = _get_config_value("ai_model")
    cache_key_val = f"{force_provider}:{force_model}:{need_vision}"

    if _cached_provider is not None and _cached_provider_force == cache_key_val:
        return _cached_provider

    info = _detect_provider_impl(force_provider, force_model, need_vision)
    _cached_provider = info
    _cached_provider_force = cache_key_val
    return info


def _detect_provider_impl(
    force_provider: str | None,
    force_model: str | None,
    need_vision: bool,
) -> ProviderInfo:
    """Internal provider detection logic."""

    # If user explicitly set a provider, respect it
    if force_provider:
        fp = force_provider.lower().strip()

        if fp in ("none", "off", "disabled"):
            return ProviderInfo("none", "", None, None, False, "AI disabled by user")

        if fp == "ollama":
            running, models = check_ollama()
            if not running:
                return ProviderInfo(
                    "none", "", None, None, False,
                    "Ollama requested but not running",
                )
            model = force_model or _pick_ollama_model(models, need_vision) or "gemma4"
            base = model.split(":")[0]
            return ProviderInfo(
                "ollama", model,
                f"{OLLAMA_BASE_URL}/v1",
                "ollama",
                base in _VISION_MODELS or model in _VISION_MODELS,
                f"Ollama ({model}, local)",
            )

        if fp == "gemini":
            api_key = load_api_key()
            if not api_key:
                return ProviderInfo(
                    "none", "", None, None, False,
                    "Gemini requested but GOOGLE_API_KEY not set",
                )
            model = force_model or _DEFAULT_MODELS["gemini"]
            return ProviderInfo(
                "gemini", model, None, api_key, True,
                f"Gemini ({model}, cloud)",
            )

        if fp == "openai":
            api_key = os.environ.get("OPENAI_API_KEY")
            if not api_key:
                return ProviderInfo(
                    "none", "", None, None, False,
                    "OpenAI requested but OPENAI_API_KEY not set",
                )
            model = force_model or _DEFAULT_MODELS["openai"]
            return ProviderInfo(
                "openai", model,
                "https://api.openai.com/v1",
                api_key, True,
                f"OpenAI ({model}, cloud)",
            )

    # Auto-detect priority: Ollama → Gemini → OpenAI
    # 1. Ollama
    running, models = check_ollama()
    if running and models:
        model = force_model or _pick_ollama_model(models, need_vision)
        if model:
            base = model.split(":")[0]
            return ProviderInfo(
                "ollama", model,
                f"{OLLAMA_BASE_URL}/v1",
                "ollama",
                base in _VISION_MODELS or model in _VISION_MODELS,
                f"Ollama ({model}, local, free)",
            )

    # 2. Gemini
    api_key = load_api_key()
    if api_key:
        model = force_model or _DEFAULT_MODELS["gemini"]
        return ProviderInfo(
            "gemini", model, None, api_key, True,
            f"Gemini ({model}, cloud)",
        )

    # 3. OpenAI
    openai_key = os.environ.get("OPENAI_API_KEY")
    if openai_key:
        model = force_model or _DEFAULT_MODELS["openai"]
        return ProviderInfo(
            "openai", model,
            "https://api.openai.com/v1",
            openai_key, True,
            f"OpenAI ({model}, cloud)",
        )

    return ProviderInfo("none", "", None, None, False, "No AI provider available")


def reset_provider_cache() -> None:
    """Reset cached provider detection (useful after config changes)."""
    global _cached_provider, _cached_provider_force
    _cached_provider = None
    _cached_provider_force = None


# ---------------------------------------------------------------------------
# Unified LLM calling
# ---------------------------------------------------------------------------


class LLMNotAvailableError(Exception):
    """Raised when no AI provider is available."""
    pass


def get_setup_message() -> str:
    """Return helpful message when no provider is found."""
    return (
        "AI features require either:\n"
        "  \u2022 Ollama (free, local): curl -fsSL https://ollama.com/install.sh | sh && ollama pull gemma4\n"
        "  \u2022 Gemini API key: export GOOGLE_API_KEY=your-key\n"
        "  \u2022 OpenAI API key: export OPENAI_API_KEY=your-key"
    )


def call_llm(
    prompt: str,
    *,
    system: str | None = None,
    images: list[bytes] | None = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
    conversation: list[dict] | None = None,
    provider_info: ProviderInfo | None = None,
) -> str | None:
    """Call the best available LLM provider.

    Auto-detects: Ollama (local) -> Gemini (cloud) -> OpenAI -> None.
    Supports text-only and multimodal (vision) requests.

    Args:
        prompt: The prompt text (used when conversation is None).
        system: Optional system prompt.
        images: Optional list of image bytes for vision requests.
        temperature: Sampling temperature.
        max_tokens: Maximum output tokens.
        conversation: Optional multi-turn messages as
            [{"role": "user"|"assistant", "content": "..."}].
        provider_info: Explicit provider to use (skips auto-detection).

    Returns:
        Response text, or None on any failure.
    """
    need_vision = images is not None and len(images) > 0
    info = provider_info or detect_provider(need_vision=need_vision)

    if info.provider == "none":
        logger.debug("No AI provider available: %s", info.display_name)
        return None

    if need_vision and not info.supports_vision:
        logger.warning(
            "Model %s may not support vision. For best results, try: ollama pull gemma4",
            info.model,
        )

    if info.provider == "gemini":
        return _call_gemini_provider(
            prompt, info, system=system, images=images,
            temperature=temperature, max_tokens=max_tokens,
            conversation=conversation,
        )
    elif info.provider in ("ollama", "openai"):
        return _call_openai_compat(
            prompt, info, system=system, images=images,
            temperature=temperature, max_tokens=max_tokens,
            conversation=conversation,
        )

    return None


def _call_gemini_provider(
    prompt: str,
    info: ProviderInfo,
    *,
    system: str | None = None,
    images: list[bytes] | None = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
    conversation: list[dict] | None = None,
) -> str | None:
    """Call Gemini via REST API (existing code path)."""
    if conversation:
        contents = [
            {
                "role": "model" if m["role"] == "assistant" else "user",
                "parts": [{"text": m["content"]}],
            }
            for m in conversation
        ]
    elif images:
        parts: list[dict] = []
        for img_data in images[:5]:
            if isinstance(img_data, bytes):
                b64 = base64.b64encode(img_data).decode()
            else:
                b64 = img_data  # already base64 string
            parts.append({
                "inline_data": {
                    "mime_type": "image/jpeg",
                    "data": b64,
                }
            })
        parts.append({"text": prompt})
        contents = [{"role": "user", "parts": parts}]
    else:
        contents = [{"role": "user", "parts": [{"text": prompt}]}]

    body: dict = {
        "contents": contents,
        "generationConfig": {
            "temperature": temperature,
            "maxOutputTokens": max_tokens,
        },
    }
    if system:
        body["system_instruction"] = {"parts": [{"text": system}]}

    url = _GEMINI_URL.format(model=info.model, key=info.api_key)
    try:
        req = urllib.request.Request(
            url, data=json.dumps(body).encode(),
            headers={"Content-Type": "application/json"},
        )
        timeout = 60 if images else 30
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read())
        candidate = data["candidates"][0]
        finish = candidate.get("finishReason", "")
        if finish not in ("STOP", "MAX_TOKENS"):
            logger.debug("Gemini finish reason: %s", finish)
            return None
        text = candidate["content"]["parts"][0]["text"]
        if finish == "MAX_TOKENS":
            logger.debug("Gemini response truncated at max_tokens")
        return text
    except Exception as e:
        logger.debug("Gemini call failed: %s", e)
        return None


def _call_openai_compat(
    prompt: str,
    info: ProviderInfo,
    *,
    system: str | None = None,
    images: list[bytes] | None = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
    conversation: list[dict] | None = None,
) -> str | None:
    """Call Ollama or OpenAI-compatible provider via the openai Python client."""
    try:
        from openai import OpenAI
    except ImportError:
        # Fall back to raw HTTP if openai package not installed
        return _call_openai_compat_raw(
            prompt, info, system=system, images=images,
            temperature=temperature, max_tokens=max_tokens,
            conversation=conversation,
        )

    client = OpenAI(
        base_url=info.base_url,
        api_key=info.api_key,
    )

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

    if conversation:
        messages.extend(conversation)
    elif images:
        content: list[dict] = []
        for img_data in images[:5]:
            if isinstance(img_data, bytes):
                b64 = base64.b64encode(img_data).decode()
            else:
                b64 = img_data
            content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
            })
        content.append({"type": "text", "text": prompt})
        messages.append({"role": "user", "content": content})
    else:
        messages.append({"role": "user", "content": prompt})

    try:
        timeout_val = 180 if images else 120
        response = client.chat.completions.create(
            model=info.model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            timeout=timeout_val,
        )
        return response.choices[0].message.content
    except Exception as e:
        logger.debug("OpenAI-compat call failed (%s): %s", info.provider, e)
        return None


def _call_openai_compat_raw(
    prompt: str,
    info: ProviderInfo,
    *,
    system: str | None = None,
    images: list[bytes] | None = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
    conversation: list[dict] | None = None,
) -> str | None:
    """Fallback: call OpenAI-compatible API via raw HTTP (no openai package)."""
    messages: list[dict] = []
    if system:
        messages.append({"role": "system", "content": system})

    if conversation:
        messages.extend(conversation)
    elif images:
        content: list[dict] = []
        for img_data in images[:5]:
            if isinstance(img_data, bytes):
                b64 = base64.b64encode(img_data).decode()
            else:
                b64 = img_data
            content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
            })
        content.append({"type": "text", "text": prompt})
        messages.append({"role": "user", "content": content})
    else:
        messages.append({"role": "user", "content": prompt})

    body = json.dumps({
        "model": info.model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }).encode()

    url = f"{info.base_url}/chat/completions"
    try:
        req = urllib.request.Request(
            url, data=body,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {info.api_key}",
            },
        )
        timeout_val = 180 if images else 120
        with urllib.request.urlopen(req, timeout=timeout_val) as resp:
            data = json.loads(resp.read())
        return data["choices"][0]["message"]["content"]
    except Exception as e:
        logger.debug("OpenAI-compat raw call failed: %s", e)
        return None


# ---------------------------------------------------------------------------
# Backward-compatible wrapper
# ---------------------------------------------------------------------------


def call_gemini(
    prompt: str,
    api_key: str,
    temperature: float = 0.0,
    max_tokens: int = 4096,
    conversation: list[dict] | None = None,
) -> str | None:
    """Call Gemini Flash-Lite and return the text response.

    Backward-compatible wrapper — new code should use call_llm() instead.
    """
    info = ProviderInfo(
        provider="gemini",
        model=_DEFAULT_MODELS["gemini"],
        base_url=None,
        api_key=api_key,
        supports_vision=True,
        display_name="Gemini (legacy)",
    )
    return _call_gemini_provider(
        prompt, info,
        temperature=temperature,
        max_tokens=max_tokens,
        conversation=conversation,
    )


# ---------------------------------------------------------------------------
# Disk cache
# ---------------------------------------------------------------------------


def cache_key(data: str) -> str:
    """Deterministic short hash for cache lookup."""
    return hashlib.md5(data.encode()).hexdigest()[:16]


def cache_get(key: str, category: str = "general") -> str | None:
    """Read a cached value, or None if miss."""
    f = CACHE_DIR / category / f"{key}.txt"
    if f.exists():
        try:
            return f.read_text()
        except OSError:
            return None
    return None


def cache_set(key: str, value: str, category: str = "general") -> None:
    """Write a value to the disk cache."""
    d = CACHE_DIR / category
    try:
        d.mkdir(parents=True, exist_ok=True)
        (d / f"{key}.txt").write_text(value)
    except OSError as e:
        logger.debug("Cache write failed (best-effort): %s", e)


# ---------------------------------------------------------------------------
# Provider display helper
# ---------------------------------------------------------------------------


def provider_status_line(info: ProviderInfo | None = None) -> str:
    """Return a one-line status message about the current AI provider."""
    if info is None:
        info = detect_provider()

    if info.provider == "ollama":
        return f"Using local AI ({info.model} via Ollama) \u2014 free, no data leaves your machine"
    elif info.provider == "gemini":
        return f"Using {info.model} (cloud)"
    elif info.provider == "openai":
        return f"Using {info.model} (OpenAI cloud)"
    else:
        return "No AI provider available"
