"""
LLM Router — Unified interface that discovers and routes to all available LLM backends.

Auto-detects which backends are reachable (local servers) or configured (API keys)
and provides a single ``chat()`` / ``quick_chat()`` surface for the rest of Brynq.

Supported backends:
  1. Ollama          — localhost:11434, native /api/chat
  2. LM Studio       — localhost:1234,  OpenAI-compatible /v1/chat/completions
  3. Claude (API)    — api.anthropic.com, requires ANTHROPIC_API_KEY
  4. Gemini (API)    — generativelanguage.googleapis.com, requires GEMINI_API_KEY

All stdlib — no external packages.

Usage:
    from brynq.llm_router import discover_backends, chat, quick_chat, get_best_backend

    # See what is available
    backends = discover_backends()

    # One-shot helper (auto-selects backend + model)
    reply = quick_chat("Explain TCP in one sentence.")

    # Explicit backend / model
    reply = chat("claude", "claude-sonnet-4-20250514",
                 [{"role": "user", "content": "Hi"}])

    # Pick the best backend for a task type
    backend, model = get_best_backend("code")
"""

from __future__ import annotations

import json
import logging
import os
import ssl
import urllib.error
import urllib.request
from typing import Any, Optional

log = logging.getLogger("brynq.llm_router")

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_OLLAMA_URL = "http://localhost:11434"
_LMSTUDIO_URL = "http://localhost:1234"
_ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
_GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta/models"

_ANTHROPIC_VERSION = "2023-06-01"
_DEFAULT_MAX_TOKENS = 2048
_HTTP_TIMEOUT = 120   # seconds — generous for large generations
_PROBE_TIMEOUT = 4    # seconds — quick reachability check

# Default models when the caller says "auto"
_CLAUDE_DEFAULT_MODEL = "claude-sonnet-4-20250514"
_GEMINI_DEFAULT_MODEL = "gemini-2.0-flash"

# Legacy aliases so callers that import ``is_available`` with old backend names
# still work (openai, vllm, llamacpp all route to OpenAI-compatible).
_DEFAULTS = {
    "ollama": _OLLAMA_URL,
    "openai": f"{_LMSTUDIO_URL}/v1",
    "lmstudio": f"{_LMSTUDIO_URL}/v1",
    "vllm": "http://localhost:8000/v1",
    "llamacpp": "http://localhost:8080/v1",
}

# ---------------------------------------------------------------------------
# Internal HTTP helpers (stdlib only)
# ---------------------------------------------------------------------------


def _build_ssl_context() -> ssl.SSLContext:
    """Return a default SSL context for HTTPS requests."""
    return ssl.create_default_context()


def _http_json(
    method: str,
    url: str,
    body: Optional[dict] = None,
    headers: Optional[dict[str, str]] = None,
    timeout: int = _HTTP_TIMEOUT,
) -> dict:
    """Make an HTTP request and return parsed JSON.

    Raises ``RuntimeError`` on HTTP or network errors.
    """
    hdrs: dict[str, str] = {
        "Content-Type": "application/json",
        "Accept": "application/json",
    }
    if headers:
        hdrs.update(headers)

    data = json.dumps(body).encode("utf-8") if body is not None else None
    req = urllib.request.Request(url, data=data, headers=hdrs, method=method)

    context = _build_ssl_context() if url.startswith("https://") else None

    try:
        with urllib.request.urlopen(req, timeout=timeout, context=context) as resp:
            raw = resp.read().decode("utf-8")
            return json.loads(raw) if raw.strip() else {}
    except urllib.error.HTTPError as exc:
        detail = ""
        try:
            detail = exc.read().decode("utf-8", errors="replace")
        except Exception:
            pass
        raise RuntimeError(f"HTTP {exc.code} {method} {url}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"Connection failed {method} {url}: {exc.reason}") from exc


def _probe_url(url: str, timeout: int = _PROBE_TIMEOUT) -> bool:
    """Return ``True`` if *url* responds to a GET within *timeout* seconds."""
    req = urllib.request.Request(url, method="GET")
    try:
        with urllib.request.urlopen(req, timeout=timeout):
            return True
    except Exception:
        return False


# ---------------------------------------------------------------------------
# Backend probes
# ---------------------------------------------------------------------------


def _probe_ollama() -> dict[str, Any]:
    """Check if Ollama is running and enumerate its models."""
    info: dict[str, Any] = {
        "name": "ollama",
        "status": "unavailable",
        "models": [],
        "url": "",
    }
    if not _probe_url(_OLLAMA_URL):
        return info

    info["url"] = "localhost:11434"
    try:
        resp = _http_json("GET", f"{_OLLAMA_URL}/api/tags", timeout=_PROBE_TIMEOUT)
        models = [
            m.get("name", "") for m in resp.get("models", []) if m.get("name")
        ]
        info["models"] = models
        info["status"] = "running"
    except RuntimeError:
        # Server responded to the root probe but /api/tags failed — still running
        info["status"] = "running"
    return info


def _probe_lmstudio() -> dict[str, Any]:
    """Check if LM Studio (or any OpenAI-compatible server) is on port 1234."""
    info: dict[str, Any] = {
        "name": "lmstudio",
        "status": "unavailable",
        "models": [],
        "url": "",
    }
    if not _probe_url(f"{_LMSTUDIO_URL}/v1/models"):
        return info

    info["url"] = "localhost:1234"
    try:
        resp = _http_json(
            "GET", f"{_LMSTUDIO_URL}/v1/models", timeout=_PROBE_TIMEOUT
        )
        models = [m.get("id", "") for m in resp.get("data", []) if m.get("id")]
        info["models"] = models
        info["status"] = "running"
    except RuntimeError:
        info["status"] = "running"
    return info


def _probe_claude() -> dict[str, Any]:
    """Check if the Anthropic API key is set."""
    info: dict[str, Any] = {
        "name": "claude",
        "status": "not_configured",
        "models": [],
        "url": "",
    }
    if os.environ.get("ANTHROPIC_API_KEY", "").strip():
        info["status"] = "configured"
        info["models"] = [_CLAUDE_DEFAULT_MODEL]
        info["url"] = "api.anthropic.com"
    return info


def _probe_gemini() -> dict[str, Any]:
    """Check if the Gemini API key is set."""
    info: dict[str, Any] = {
        "name": "gemini",
        "status": "not_configured",
        "models": [],
        "url": "",
    }
    if os.environ.get("GEMINI_API_KEY", "").strip():
        info["status"] = "configured"
        info["models"] = [_GEMINI_DEFAULT_MODEL]
        info["url"] = "generativelanguage.googleapis.com"
    return info


# ---------------------------------------------------------------------------
# Public: discover_backends
# ---------------------------------------------------------------------------


def discover_backends() -> list[dict]:
    """Return all known backends with their availability and loaded models.

    Each entry has the shape::

        {
            "name":   "ollama" | "lmstudio" | "claude" | "gemini",
            "status": "running" | "configured" | "not_configured" | "unavailable",
            "models": ["model-name", ...],
            "url":    "host:port or domain",
        }

    ``"running"`` means a local server was reachable.
    ``"configured"`` means the required API-key env-var is set.
    ``"not_configured"`` means the env-var is missing.
    ``"unavailable"`` means the local server did not respond.
    """
    return [
        _probe_ollama(),
        _probe_lmstudio(),
        _probe_claude(),
        _probe_gemini(),
    ]


# ---------------------------------------------------------------------------
# Backend chat implementations
# ---------------------------------------------------------------------------


def _chat_ollama(model: str, messages: list[dict], **kwargs: Any) -> str:
    """Send a chat request through the Ollama native API."""
    payload: dict[str, Any] = {
        "model": model,
        "messages": messages,
        "stream": False,
    }
    if "temperature" in kwargs:
        payload["options"] = {"temperature": kwargs["temperature"]}

    resp = _http_json("POST", f"{_OLLAMA_URL}/api/chat", body=payload)
    return resp.get("message", {}).get("content", "")


def _chat_lmstudio(model: str, messages: list[dict], **kwargs: Any) -> str:
    """Send a chat request to an OpenAI-compatible server on port 1234."""
    payload: dict[str, Any] = {
        "model": model,
        "messages": messages,
    }
    if "temperature" in kwargs:
        payload["temperature"] = kwargs["temperature"]
    if "max_tokens" in kwargs:
        payload["max_tokens"] = kwargs["max_tokens"]

    resp = _http_json(
        "POST",
        f"{_LMSTUDIO_URL}/v1/chat/completions",
        body=payload,
        headers={"Authorization": "Bearer not-needed"},
    )
    choices = resp.get("choices", [])
    if not choices:
        return ""
    return choices[0].get("message", {}).get("content", "")


def _chat_claude(model: str, messages: list[dict], **kwargs: Any) -> str:
    """Send a chat request to the Anthropic Messages API.

    Converts from the standard ``[{"role": ..., "content": ...}]`` format.
    The Anthropic API requires that messages alternate user/assistant and
    that system is a top-level parameter, not inside the messages list.

    If ``api_key`` is passed in kwargs, it takes priority over the env var.
    """
    api_key = kwargs.pop("api_key", None) or os.environ.get("ANTHROPIC_API_KEY", "").strip()
    if not api_key:
        raise RuntimeError("ANTHROPIC_API_KEY environment variable is not set")

    # Separate system prompt from conversation messages
    system_text: Optional[str] = None
    api_messages: list[dict[str, str]] = []
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        if role == "system":
            system_text = (
                content
                if system_text is None
                else system_text + "\n\n" + content
            )
        else:
            # Anything other than user/assistant becomes user
            if role not in ("user", "assistant"):
                role = "user"
            api_messages.append({"role": role, "content": content})

    if not api_messages:
        api_messages = [{"role": "user", "content": "Hello."}]

    # Merge consecutive same-role messages (Anthropic rejects them)
    merged: list[dict[str, str]] = []
    for msg in api_messages:
        if merged and merged[-1]["role"] == msg["role"]:
            merged[-1]["content"] += "\n\n" + msg["content"]
        else:
            merged.append(dict(msg))

    # Must start with a user message
    if merged[0]["role"] != "user":
        merged.insert(0, {"role": "user", "content": "(continued)"})

    body: dict[str, Any] = {
        "model": model,
        "max_tokens": kwargs.get("max_tokens", _DEFAULT_MAX_TOKENS),
        "messages": merged,
    }
    if system_text:
        body["system"] = system_text
    if "temperature" in kwargs:
        body["temperature"] = kwargs["temperature"]

    resp = _http_json(
        "POST",
        _ANTHROPIC_API_URL,
        body=body,
        headers={
            "x-api-key": api_key,
            "anthropic-version": _ANTHROPIC_VERSION,
        },
    )

    # Response shape: {"content": [{"type": "text", "text": "..."}], ...}
    parts: list[str] = []
    for block in resp.get("content", []):
        if block.get("type") == "text":
            parts.append(block.get("text", ""))
    return "\n".join(parts)


def _chat_gemini(model: str, messages: list[dict], **kwargs: Any) -> str:
    """Send a chat request to the Google Gemini / Generative AI API.

    Converts the standard messages list into Gemini's ``contents`` schema.

    If ``api_key`` is passed in kwargs, it takes priority over the env var.
    """
    api_key = kwargs.pop("api_key", None) or os.environ.get("GEMINI_API_KEY", "").strip()
    if not api_key:
        raise RuntimeError("GEMINI_API_KEY environment variable is not set")

    # Gemini uses "user" and "model" (not "assistant")
    contents: list[dict[str, Any]] = []
    system_parts: list[str] = []

    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        if role == "system":
            system_parts.append(content)
            continue
        gemini_role = "model" if role == "assistant" else "user"
        contents.append({
            "role": gemini_role,
            "parts": [{"text": content}],
        })

    # Prepend system instructions into the first user turn
    if system_parts and contents:
        system_text = "\n\n".join(system_parts)
        if contents[0]["role"] == "user":
            existing = contents[0]["parts"][0].get("text", "")
            contents[0]["parts"][0]["text"] = system_text + "\n\n" + existing
        else:
            contents.insert(0, {
                "role": "user",
                "parts": [{"text": system_text}],
            })

    if not contents:
        contents = [{"role": "user", "parts": [{"text": "Hello."}]}]

    # Merge consecutive same-role entries
    merged: list[dict[str, Any]] = []
    for c in contents:
        if merged and merged[-1]["role"] == c["role"]:
            merged[-1]["parts"].extend(c["parts"])
        else:
            merged.append(dict(c))

    # Must start with user
    if merged[0]["role"] != "user":
        merged.insert(0, {"role": "user", "parts": [{"text": "(continued)"}]})

    body: dict[str, Any] = {"contents": merged}

    gen_config: dict[str, Any] = {}
    if "temperature" in kwargs:
        gen_config["temperature"] = kwargs["temperature"]
    if "max_tokens" in kwargs:
        gen_config["maxOutputTokens"] = kwargs["max_tokens"]
    if gen_config:
        body["generationConfig"] = gen_config

    url = f"{_GEMINI_API_BASE}/{model}:generateContent?key={api_key}"
    resp = _http_json("POST", url, body=body)

    # Response: {"candidates": [{"content": {"parts": [{"text": "..."}]}}]}
    candidates = resp.get("candidates", [])
    if not candidates:
        return ""
    text_parts: list[str] = []
    for p in candidates[0].get("content", {}).get("parts", []):
        if "text" in p:
            text_parts.append(p["text"])
    return "\n".join(text_parts)


# ---------------------------------------------------------------------------
# Dispatch table
# ---------------------------------------------------------------------------

_BACKEND_DISPATCH: dict[str, Any] = {
    "ollama": _chat_ollama,
    "lmstudio": _chat_lmstudio,
    "claude": _chat_claude,
    "gemini": _chat_gemini,
}

_BACKEND_STATUS_USABLE = {"running", "configured"}

# ---------------------------------------------------------------------------
# Public: chat
# ---------------------------------------------------------------------------


def chat(
    backend: str,
    model: str,
    messages: list[dict],
    *,
    api_key: Optional[str] = None,
    **kwargs: Any,
) -> str:
    """Route a chat request to the specified backend.

    Args:
        backend: One of ``"ollama"``, ``"lmstudio"``, ``"claude"``, ``"gemini"``,
                 or ``"auto"`` to pick the first available backend.
        model:   Model name (e.g. ``"llama3"``, ``"claude-sonnet-4-20250514"``).
                 Pass ``"auto"`` to use the backend's default or first loaded model.
        messages: Conversation in the standard chat format::

                     [
                         {"role": "system",    "content": "..."},
                         {"role": "user",      "content": "..."},
                         {"role": "assistant", "content": "..."},
                     ]

        api_key: Optional API key override (BYOK). If provided, it is used
                 instead of the environment variable for cloud backends
                 (claude, gemini). Ignored for local backends (ollama, lmstudio).
        **kwargs: Forwarded to the backend (e.g. ``temperature``, ``max_tokens``).

    Returns:
        The assistant's text response.

    Raises:
        RuntimeError: If the backend is unknown, unreachable, or missing its
                      API key.
    """
    if backend == "auto":
        backend, auto_model = _pick_first_available()
        if model == "auto":
            model = auto_model

    if model == "auto":
        model = _default_model_for(backend)

    handler = _BACKEND_DISPATCH.get(backend)
    if handler is None:
        raise RuntimeError(
            f"Unknown backend '{backend}'. "
            f"Supported: {', '.join(sorted(_BACKEND_DISPATCH))}"
        )

    # Pass api_key through kwargs for cloud backends that support it
    if api_key and backend in ("claude", "gemini"):
        kwargs["api_key"] = api_key

    return handler(model, messages, **kwargs)


# ---------------------------------------------------------------------------
# Public: quick_chat
# ---------------------------------------------------------------------------


def quick_chat(
    prompt: str,
    backend: str = "auto",
    model: str = "auto",
    **kwargs: Any,
) -> str:
    """Simplified one-shot chat.  Auto-selects backend and model if unspecified.

    Args:
        prompt:  The user's message.
        backend: Backend name or ``"auto"``.
        model:   Model name or ``"auto"``.
        **kwargs: Forwarded to ``chat()`` (temperature, max_tokens, etc.).

    Returns:
        The assistant's text response.

    Raises:
        RuntimeError: If no backend is available.
    """
    messages = [{"role": "user", "content": prompt}]
    return chat(backend=backend, model=model, messages=messages, **kwargs)


# ---------------------------------------------------------------------------
# Public: get_best_backend
# ---------------------------------------------------------------------------


def get_best_backend(task_type: str = "general") -> tuple[str, str]:
    """Return the best ``(backend_name, model_name)`` pair for a task type.

    Task types:
        ``"code"``    -- Prefer Claude, then capable local models.
        ``"fast"``    -- Prefer the smallest available local model.
        ``"quality"`` -- Prefer Claude, then Gemini API.
        ``"general"`` -- First available backend, any model.

    Raises:
        RuntimeError: If no backend is available.
    """
    backends = {b["name"]: b for b in discover_backends()}
    usable = {
        name: info
        for name, info in backends.items()
        if info["status"] in _BACKEND_STATUS_USABLE
    }

    if not usable:
        raise RuntimeError("No LLM backends available")

    def _first_model(name: str) -> str:
        models = usable[name]["models"]
        return models[0] if models else _default_model_for(name)

    if task_type == "code":
        for name in ("claude", "ollama", "lmstudio", "gemini"):
            if name in usable:
                return name, _first_model(name)

    elif task_type == "fast":
        # Prefer local backends for lowest latency
        for name in ("ollama", "lmstudio"):
            if name in usable:
                models = usable[name]["models"]
                if models:
                    # Shortest name as a rough proxy for smallest model
                    return name, min(models, key=len)
                return name, _default_model_for(name)
        for name in ("gemini", "claude"):
            if name in usable:
                return name, _first_model(name)

    elif task_type == "quality":
        for name in ("claude", "gemini", "ollama", "lmstudio"):
            if name in usable:
                return name, _first_model(name)

    # "general" or unrecognised task_type — first available
    for name in ("ollama", "lmstudio", "claude", "gemini"):
        if name in usable:
            return name, _first_model(name)

    # Fallback (should not reach here given the guard above)
    name = next(iter(usable))
    return name, _first_model(name)


# ---------------------------------------------------------------------------
# Legacy API — kept for backward compatibility
# ---------------------------------------------------------------------------


def is_available(backend: str, base_url: Optional[str] = None) -> bool:
    """Check whether a backend is reachable.

    Retained for backward compatibility.  New code should prefer
    ``discover_backends()`` which returns richer status info.

    Args:
        backend: Backend name (``"ollama"``, ``"lmstudio"``, ``"openai"``,
                 ``"claude"``, ``"gemini"``, ``"vllm"``, ``"llamacpp"``).
        base_url: Override the default endpoint URL.

    Returns:
        ``True`` if the backend is reachable or its API key is set.
    """
    if backend in ("claude", "gemini"):
        return True  # Bypass API key check since the desktop CLI handles its own auth

    url = (
        base_url.rstrip("/")
        if base_url
        else _DEFAULTS.get(backend, _DEFAULTS["openai"])
    )
    try:
        if backend == "ollama":
            # The root endpoint for Ollama returns "Ollama is running" which is plain text, not JSON.
            # We should ping /api/tags which returns valid JSON to check availability.
            _http_json("GET", f"{url}/api/tags", timeout=_PROBE_TIMEOUT)
        else:
            _http_json("GET", f"{url}/models", timeout=_PROBE_TIMEOUT)
        return True
    except RuntimeError:
        return False


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _pick_first_available() -> tuple[str, str]:
    """Return ``(backend, model)`` for the first usable backend.

    Raises ``RuntimeError`` if nothing is available.
    """
    for probe_fn, name in [
        (_probe_ollama, "ollama"),
        (_probe_lmstudio, "lmstudio"),
        (_probe_claude, "claude"),
        (_probe_gemini, "gemini"),
    ]:
        info = probe_fn()
        if info["status"] in _BACKEND_STATUS_USABLE:
            models = info.get("models", [])
            model = models[0] if models else _default_model_for(name)
            return name, model

    raise RuntimeError(
        "No LLM backends available. Start Ollama or LM Studio, "
        "or set ANTHROPIC_API_KEY / GEMINI_API_KEY."
    )


def _default_model_for(backend: str) -> str:
    """Return a sensible default model name for *backend*."""
    return {
        "ollama": "llama3",
        "lmstudio": "local-model",
        "claude": _CLAUDE_DEFAULT_MODEL,
        "gemini": _GEMINI_DEFAULT_MODEL,
    }.get(backend, "default")
