"""Vector storage and retrieval with USearch."""

import numpy as np
from usearch.index import Index


def embedding_store_usearch(path: str, ids: list, embeddings: list, dim: int = 384) -> int:
    """Crea o sobreescribe un indice USearch con los embeddings dados.

    Args:
        path: Path del archivo .usearch para persistir el indice.
        ids: Lista de IDs (int) para cada embedding.
        embeddings: Lista de arrays numpy float32.
        dim: Dimension de los embeddings (default 384 para e5-small).

    Returns:
        Numero de embeddings insertados.

    Raises:
        OSError: Si no se puede escribir en el path.
    """
    index = Index(ndim=dim, metric="ip", dtype="f32")
    keys = np.array(ids, dtype=np.uint64)
    vecs = np.array(embeddings, dtype=np.float32)
    index.add(keys, vecs)
    index.save(path)
    return len(ids)


def embedding_search_usearch(path: str, query_embedding: list, k: int = 10, dim: int = 384) -> list:
    """Busca los k vecinos mas cercanos en un indice USearch persistido.

    Args:
        path: Path del archivo .usearch.
        query_embedding: Array numpy float32 del query.
        k: Numero de resultados a retornar.
        dim: Dimension de los embeddings.

    Returns:
        Lista de dicts con 'key' y 'distance' ordenados por cercania.

    Raises:
        OSError: Si el archivo no existe.
    """
    index = Index(ndim=dim, metric="ip", dtype="f32")
    index.load(path)
    q = np.asarray(query_embedding, dtype=np.float32)
    results = index.search(q, k)
    keys = np.atleast_1d(results.keys)
    distances = np.atleast_1d(results.distances)
    return [{"key": int(keys[i]), "distance": float(distances[i])} for i in range(len(keys))]