"""Embedding model management — save, load, and encode with multilingual-e5-small."""

import os

from sentence_transformers import SentenceTransformer


def embedding_save_model(model_id: str, path: str) -> str:
    """Descarga modelo de HuggingFace y lo guarda en path local.

    Args:
        model_id: ID del modelo en HuggingFace (ej: "intfloat/multilingual-e5-small").
        path: Directorio destino para guardar el modelo.

    Returns:
        Path absoluto donde se guardo el modelo.

    Raises:
        OSError: Si no se puede escribir en el path.
        Exception: Si el modelo no existe en HuggingFace.
    """
    os.makedirs(path, exist_ok=True)
    model = SentenceTransformer(model_id)
    model.save(path)
    return os.path.abspath(path)


def embedding_load_model(path: str) -> SentenceTransformer:
    """Carga modelo de embeddings desde path local.

    Args:
        path: Directorio con el modelo guardado por embedding_save_model.

    Returns:
        Instancia de SentenceTransformer lista para encode.

    Raises:
        OSError: Si el path no existe o no contiene un modelo valido.
    """
    return SentenceTransformer(path)


def embedding_encode(model: SentenceTransformer, texts: list, mode: str = "document") -> list:
    """Genera embeddings normalizados para una lista de textos.

    Aplica automaticamente los prefijos requeridos por modelos e5:
    - mode="document" -> "passage: " prefix
    - mode="query"    -> "query: " prefix

    Args:
        model: Modelo cargado con embedding_load_model.
        texts: Lista de strings a codificar.
        mode: "document" para indexar, "query" para buscar.

    Returns:
        Lista de arrays numpy float32 normalizados (dim depende del modelo).

    Raises:
        ValueError: Si mode no es "document" ni "query".
    """
    if mode not in ("document", "query"):
        raise ValueError(f"mode must be 'document' or 'query', got '{mode}'")

    prefix = "passage: " if mode == "document" else "query: "
    prefixed = [f"{prefix}{t}" for t in texts]
    embeddings = model.encode(prefixed, normalize_embeddings=True, show_progress_bar=False)
    return embeddings