"""The model table the LLM gateway routes with.

``build_model_list`` turns the TinyAgentOS backend list into the routing
table (model name -> backend) that ``llm_gateway.resolve`` and
``llm_gateway.embeddings`` read. Its entries keep the LiteLLM-shaped
``litellm_params`` keys they had when this table also fed the LiteLLM proxy
config; that writer and the proxy are gone (LiteLLM removal stage 2b-2a).
The module keeps its name until the cosmetic rename (stage 2b-2b).
"""
from __future__ import annotations

import logging

import httpx

from tinyagentos.providers import (
    ALL_TYPES,
    BACKEND_TYPE_MAP,
    CHAT_BACKEND_TYPE_MAP,
    CLOUD_TYPES as CLOUD_BACKEND_TYPES,
    NEEDS_API_BASE_TYPES,
    inference_api_base,
)

logger = logging.getLogger(__name__)

# Canonical alias the deployer injects into agent containers as
# TAOS_EMBEDDING_MODEL. Agents that want an embedding call this name and
# LiteLLM routes it to whatever concrete embedding model the host has.
# See docs/design/framework-agnostic-runtime.md.
EMBEDDING_ALIAS = "taos-embedding-default"

def _is_embedding_model(name: str) -> bool:
    """Classify a model name as embedding vs chat.

    Heuristic match against the common embedding-model families. LiteLLM
    has no reliable way to ask a backend "is this model an embedder?" so
    we infer from the slug. Known families: anything containing ``embed``
    (nomic-embed, qwen3-embedding, mxbai-embed-large), plus the BGE, GTE,
    E5, and arctic-embed families whose canonical names don't always
    include the word.

    Rerankers include ``rerank`` in the slug and are always excluded —
    they're a different endpoint shape that LiteLLM doesn't front today.
    """
    n = name.lower()
    if "rerank" in n:
        return False
    if "embed" in n:
        return True
    # Known embedding families where the slug doesn't include "embed"
    # verbatim. Match on dash-separated prefixes so "bge-something" hits
    # but a hypothetical "bgem-chat" doesn't.
    for prefix in ("bge-", "gte-", "e5-", "arctic-"):
        if n.startswith(prefix):
            return True
    return False


def _discover_ollama_models(url: str, timeout: float = 2.0) -> list[str]:
    """Probe an ollama-compatible backend for its loaded model names.

    Returns ``[]`` on any failure (backend down, network error, schema
    drift) — the caller treats an empty list as "no models to auto-wire"
    and falls back to the ``default`` alias.
    """
    try:
        resp = httpx.get(f"{url.rstrip('/')}/api/tags", timeout=timeout)
        if resp.status_code != 200:
            return []
        data = resp.json()
        return [m.get("name", "") for m in data.get("models", []) if m.get("name")]
    except Exception as exc:
        logger.debug("ollama model probe at %s failed: %s", url, exc)
        return []


def _discover_openai_compatible_models(url: str, timeout: float = 2.0, api_key: str | None = None) -> list[str]:
    """Probe an OpenAI-compatible backend (llama-cpp, lmstudio) for its loaded model names.

    Returns ``[]`` on any failure (backend down, network error, schema
    drift) — the caller treats an empty list as "no models to auto-wire"
    and falls back to the ``default`` alias.

    OpenAI /models endpoint returns: {"data": [{"id": "model-name"}, ...]}

    If ``api_key`` is provided, an Authorization header is sent so
    key-protected servers return models instead of 401.
    """
    try:
        headers = {"Authorization": f"Bearer {api_key}"} if api_key else None
        resp = httpx.get(f"{url.rstrip('/')}/models", timeout=timeout, headers=headers)
        if resp.status_code != 200:
            return []
        data = resp.json()
        return [m.get("id", "") for m in data.get("data", []) if m.get("id")]
    except Exception as exc:
        logger.debug("openai-compatible model probe at %s failed: %s", url, exc)
        return []


async def _discover_ollama_backends_concurrent(
    backends: list[dict], timeout: float = 2.0
) -> dict[str, list[str]]:
    """Probe all ollama/rkllama backends concurrently.

    Returns a mapping of URL -> model name list. Each probe runs in a
    thread via ``asyncio.to_thread`` so blocking httpx calls do not
    stall the event loop.
    """
    import asyncio

    ollama_urls = [
        b.get("url", "").rstrip("/")
        for b in backends
        if b.get("type", "ollama") in ("ollama", "rkllama", "hailo-ollama") and b.get("url")
    ]
    if not ollama_urls:
        return {}
    results = await asyncio.gather(
        *(asyncio.to_thread(_discover_ollama_models, url, timeout) for url in ollama_urls),
        return_exceptions=True,
    )
    return {
        url: (res if isinstance(res, list) else [])
        for url, res in zip(ollama_urls, results)
    }


async def _discover_openai_compatible_backends_concurrent(
    backends: list[dict], timeout: float = 2.0, keys: dict[str, str | None] | None = None
) -> dict[str, list[str]]:
    """Probe all OpenAI-compatible backends (llama-cpp, lmstudio) concurrently.

    Mirrors ``_discover_ollama_backends_concurrent`` exactly, but targets
    backends whose type is ``llama-cpp``, or ``lmstudio`` when ``lmstudio`` is
    in ``ALL_TYPES`` (the same condition ``build_model_list`` uses). Each probe
    runs ``_discover_openai_compatible_models`` in a thread via
    ``asyncio.to_thread`` so blocking httpx calls do not stall the event loop.

    If ``keys`` is provided, it maps stripped URL -> resolved API key. The key
    for each backend's URL is passed as ``api_key`` to the probe so
    key-protected servers return models instead of 401.
    """
    import asyncio

    probes = [
        (b.get("url", "").rstrip("/"), b.get("type"))
        for b in backends
        if (b.get("type") == "llama-cpp"
            or (b.get("type") == "lmstudio" and "lmstudio" in ALL_TYPES))
        and b.get("url")
    ]
    if not probes:
        return {}
    key_map = keys or {}
    results = await asyncio.gather(
        *(
            asyncio.to_thread(
                _discover_openai_compatible_models,
                inference_api_base(backend_type, url),
                timeout,
                api_key=key_map.get(url),
            )
            for url, backend_type in probes
        ),
        return_exceptions=True,
    )
    return {
        url: (res if isinstance(res, list) else [])
        for (url, _), res in zip(probes, results)
    }


def _local_backend_models_from_registry(
    backend: dict, registry
) -> list[str]:
    """Find every installed model that can run on this local backend.

    auto_register_from_manifest names backends ``local-<service-id>`` (e.g.
    ``local-rk-llama-cpp`` for the rk-llama.cpp service). Each model
    manifest declares ``variants[].requires.backends[].id`` listing the
    runtime ids it supports. Cross-reference the two: every installed
    model whose variants point at this backend's service id is a model
    we can route requests to via this backend's URL.

    Returns a deduplicated list of manifest ids — the LiteLLM model_name
    aliases we want to create per local backend.
    """
    if registry is None:
        return []

    name = backend.get("name", "")
    if not name.startswith("local-"):
        return []
    service_id = name[len("local-"):]
    if not service_id:
        return []

    try:
        installed_rows = registry.list_installed()
    except Exception:  # noqa: BLE001
        return []
    installed_ids = {row.get("id") for row in installed_rows if row.get("id")}

    matched: list[str] = []
    for manifest_id in sorted(installed_ids):
        manifest = registry.get(manifest_id) if hasattr(registry, "get") else None
        if not manifest or getattr(manifest, "type", None) != "model":
            continue
        variants = getattr(manifest, "variants", None) or []
        for v in variants:
            if not isinstance(v, dict):
                continue
            for req in (v.get("requires", {}) or {}).get("backends", []) or []:
                if isinstance(req, dict) and req.get("id") == service_id:
                    matched.append(manifest_id)
                    break
            else:
                continue
            break
    return matched


def build_model_list(
    backends: list[dict],
    default_model: str = "default",
    *,
    registry=None,
    discovered: dict[str, list[str]] | None = None,
    discover: bool = True,
) -> list[dict]:
    """Build the routing table: every model name taOS serves, and where it goes.

    This is the ONE source of routing truth. ``generate_litellm_config`` wraps
    it for the LiteLLM proxy, and the in-process gateway
    (``tinyagentos.llm_gateway.resolve``) reads it per request, so a model
    name can never route to one backend through LiteLLM and another through
    the gateway. Pure apart from the optional ollama probe.

    ``discover=False`` skips the ``/api/tags`` probe for any ollama backend
    not already in ``discovered``: only the probe-derived embedding entries
    are then absent. The gateway uses it so a chat request never waits on a
    2s probe per ollama backend.

    Emits two kinds of model_list entries:

    - Chat entries under the ``default`` alias, routing through each
      ollama-compatible backend. Frameworks using the OpenAI SDK against
      ``OPENAI_BASE_URL`` hit these with ``model="default"``.
    - Embedding entries discovered by probing each ollama-compatible
      backend's ``/api/tags``. The first embedding model found is also
      aliased as ``taos-embedding-default`` so the deployer can inject
      a stable model name into every agent container via
      ``TAOS_EMBEDDING_MODEL``. Entries carry
      ``model_info.mode: embedding`` so LiteLLM routes
      ``/v1/embeddings`` requests to them instead of chat.
    """
    model_list = []
    sorted_backends = sorted(backends, key=lambda b: b.get("priority", 99))
    aliased_embedding_claimed = False

    for backend in sorted_backends:
        backend_type = backend.get("type", "ollama")
        # Loudly flag cloud-type entries missing url/models so the next
        # round of silent drops (the kilocode regression) is visible in
        # logs instead of surfacing only as a broken agent much later.
        if backend_type in CLOUD_BACKEND_TYPES:
            if not backend.get("url") or not backend.get("models"):
                logger.warning(
                    "backend %s skipped — missing url or models (type=%s)",
                    backend.get("name"),
                    backend_type,
                )
        prefix = CHAT_BACKEND_TYPE_MAP.get(backend_type, "openai")
        url = backend.get("url", "").rstrip("/")
        model_name = backend.get("model", "default")

        litellm_params = {
            "model": f"{prefix}/{model_name}",
        }

        # Set api_base for local/self-hosted backends and openai-compatible.
        # ``inference_api_base`` adds any version prefix the backend's OpenAI
        # surface requires (LM Studio serves under /v1); the worker's probe URL
        # stays unversioned.
        api_base = inference_api_base(backend_type, url)
        if backend_type in NEEDS_API_BASE_TYPES:
            litellm_params["api_base"] = api_base

        # API key from secrets reference
        if backend.get("api_key_secret"):
            litellm_params["api_key"] = f"os.environ/{backend['api_key_secret']}"
        elif backend.get("api_key"):
            litellm_params["api_key"] = backend["api_key"]

        # For cloud backends with a declared models list, register each model
        # by its exact id so agents can route requests to a specific model.
        if backend_type in CLOUD_BACKEND_TYPES:
            declared_models: list[str] = []
            for m in backend.get("models") or []:
                if isinstance(m, dict):
                    mid = m.get("id") or m.get("name") or ""
                else:
                    mid = str(m)
                if mid:
                    declared_models.append(mid)

            for model_id in declared_models:
                per_model_params: dict = {
                    "model": f"{prefix}/{model_id}",
                }
                # kilocode isn't a native LiteLLM provider — must set api_base
                if backend_type == "kilocode":
                    per_model_params["api_base"] = api_base
                elif url and backend_type not in ("openai",):
                    # For openrouter and other pass-through types, set api_base
                    # when an explicit url is provided
                    per_model_params["api_base"] = api_base
                # API key resolution
                if backend.get("api_key_secret"):
                    per_model_params["api_key"] = f"os.environ/{backend['api_key_secret']}"
                elif backend.get("api_key"):
                    per_model_params["api_key"] = backend["api_key"]
                model_list.append({
                    "model_name": model_id,
                    "litellm_params": per_model_params,
                    "metadata": {
                        "priority": backend.get("priority", 99),
                        "backend_name": backend.get("name", ""),
                        "backend_type": backend_type,
                    },
                })

        model_list.append({
            "model_name": default_model,
            "litellm_params": litellm_params,
            "metadata": {
                "priority": backend.get("priority", 99),
                "backend_name": backend.get("name", ""),
                "backend_type": backend_type,
            },
        })

        # For local backends (name "local-<service-id>" from
        # auto_register_from_manifest), also register every installed
        # model that targets this backend as its own LiteLLM model_name.
        # Independent of cloud/non-cloud: openai-compatible is in
        # CLOUD_BACKEND_TYPES even when it's actually a local llama-server
        # (e.g. rk-llama.cpp). Without this, an agent picker that says
        # "use gemma-4-e2b-gguf" → LiteLLM call with that model_name →
        # 400 because no such alias is registered.
        backend_name = backend.get("name", "")
        if backend_name.startswith("local-") and url:
            for manifest_id in _local_backend_models_from_registry(backend, registry):
                # Skip if an EXACT duplicate for this same backend already
                # exists (same model_name AND same api_base).  A different
                # backend serving the same manifest_id is a distinct entry
                # (multi-backend failover) and must NOT be skipped.
                if any(
                    e.get("model_name") == manifest_id
                    and e.get("litellm_params", {}).get("api_base") == api_base
                    for e in model_list
                ):
                    continue
                per_model_params: dict = {
                    "model": f"{prefix}/{manifest_id}",
                    "api_base": api_base,
                }
                if backend.get("api_key_secret"):
                    per_model_params["api_key"] = f"os.environ/{backend['api_key_secret']}"
                elif backend.get("api_key"):
                    per_model_params["api_key"] = backend["api_key"]
                model_list.append({
                    "model_name": manifest_id,
                    "litellm_params": per_model_params,
                    "metadata": {
                        "priority": backend.get("priority", 99),
                        "backend_name": backend_name,
                        "backend_type": backend_type,
                        "source": "local-installed",
                    },
                })

        # Auto-discover embedding models on ollama-compatible backends and
        # register each as its own LiteLLM entry. The first embedding model
        # found across all backends also claims the stable
        # ``taos-embedding-default`` alias so containers have one name to
        # inject regardless of which rkllama box holds the model.
        if backend_type in ("ollama", "rkllama", "hailo-ollama"):
            _probed = (discovered or {}).get(url)
            if _probed is None:
                _probed = _discover_ollama_models(url) if discover else []
            if _probed:
                for discovered_name in _probed:
                    if not _is_embedding_model(discovered_name):
                        continue
                    embed_params = {
                        "model": f"ollama/{discovered_name}",
                        "api_base": url,
                    }
                    model_list.append({
                        "model_name": discovered_name,
                        "litellm_params": embed_params,
                        "model_info": {"mode": "embedding"},
                        "metadata": {
                            "priority": backend.get("priority", 99),
                            "backend_name": backend.get("name", ""),
                            "backend_type": backend_type,
                        },
                    })
                    if not aliased_embedding_claimed:
                        model_list.append({
                            "model_name": EMBEDDING_ALIAS,
                            "litellm_params": dict(embed_params),
                            "model_info": {"mode": "embedding"},
                            "metadata": {
                                "priority": backend.get("priority", 99),
                                "backend_name": backend.get("name", ""),
                                "backend_type": backend_type,
                                "aliases": discovered_name,
                            },
                        })
                        aliased_embedding_claimed = True
        elif backend_type == "llama-cpp" or (backend_type == "lmstudio" and "lmstudio" in ALL_TYPES):
            _probed = (discovered or {}).get(url)
            if _probed is None:
                # Authenticated probing (with api_key/api_key_secret) happens
                # on the gateway request path via _discovered(); the config
                # writer has no secrets store and calls this keyless probe.
                _probed = _discover_openai_compatible_models(inference_api_base(backend_type, url)) if discover else []
            if _probed:
                embed_prefix = BACKEND_TYPE_MAP.get(backend_type, "openai")
                for discovered_name in _probed:
                    if not _is_embedding_model(discovered_name):
                        continue
                    embed_params = {
                        "model": f"{embed_prefix}/{discovered_name}",
                        "api_base": inference_api_base(backend_type, url),
                    }
                    # API key from secrets reference (same precedence as the
                    # chat entry: api_key_secret first, else api_key), so a
                    # key-protected server accepts the embedding call.
                    if backend.get("api_key_secret"):
                        embed_params["api_key"] = f"os.environ/{backend['api_key_secret']}"
                    elif backend.get("api_key"):
                        embed_params["api_key"] = backend["api_key"]
                    model_list.append({
                        "model_name": discovered_name,
                        "litellm_params": embed_params,
                        "model_info": {"mode": "embedding"},
                        "metadata": {
                            "priority": backend.get("priority", 99),
                            "backend_name": backend.get("name", ""),
                            "backend_type": backend_type,
                        },
                    })
                    if not aliased_embedding_claimed:
                        model_list.append({
                            "model_name": EMBEDDING_ALIAS,
                            "litellm_params": dict(embed_params),
                            "model_info": {"mode": "embedding"},
                            "metadata": {
                                "priority": backend.get("priority", 99),
                                "backend_name": backend.get("name", ""),
                                "backend_type": backend_type,
                                "aliases": discovered_name,
                            },
                        })
                        aliased_embedding_claimed = True

    return model_list