"""Embedding model management — download, load, and inference via ONNX."""

import click
from pathlib import Path

_model_instance = None
_model_name = None


def load_model(model_name: str):
    """Load a sentence-transformers model with ONNX backend. Caches in-process."""
    global _model_instance, _model_name
    if _model_instance is not None and _model_name == model_name:
        return _model_instance

    from sentence_transformers import SentenceTransformer

    click.echo(f"Loading model '{model_name}'...")
    try:
        _model_instance = SentenceTransformer(model_name, backend="onnx")
    except Exception:
        # Fallback: some models may not have pre-exported ONNX. Let sentence-transformers export.
        click.echo("Optimising model for ONNX inference (one-time)...")
        _model_instance = SentenceTransformer(model_name, backend="onnx")

    _model_name = model_name
    return _model_instance


def get_model_dim(model_name: str) -> int:
    """Get the embedding dimension for a model."""
    model = load_model(model_name)
    return model.get_sentence_embedding_dimension()


def embed_texts(model_name: str, texts: list[str],
                prefix: str = "", show_progress: bool = False) -> list[list[float]]:
    """Embed a list of texts, returning float vectors."""
    model = load_model(model_name)
    if prefix:
        texts = [prefix + t for t in texts]
    embeddings = model.encode(texts, show_progress_bar=show_progress, convert_to_numpy=True)
    return [e.tolist() for e in embeddings]


def download_model(model_name: str):
    """Pre-download a model (for kb init)."""
    click.echo(f"Downloading embedding model '{model_name}'...")
    load_model(model_name)
    click.echo("Embedding model ready.")


def check_model_binding(conn, cfg: dict):
    """Verify the loaded model matches what the DB expects. Raises on mismatch."""
    from kb_search.database import get_db_config

    db_model = get_db_config(conn, "model_name")
    if db_model is None:
        return  # Not yet initialised

    config_model = cfg["embedding"]["model"]
    if db_model != config_model:
        db_dim = get_db_config(conn, "embedding_dim", "?")
        raise click.ClickException(
            f"Model mismatch: DB uses '{db_model}' ({db_dim} dim) but config specifies "
            f"'{config_model}'. Run `kb reindex --model {config_model}` to switch models."
        )
