"""Embedding index manager for LociDB.

Handles embedding generation via sentence-transformers and wraps the
MemoryStore's vector search.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, List, Optional

if TYPE_CHECKING:
    from sentence_transformers import SentenceTransformer

_model_cache: dict[str, "SentenceTransformer"] = {}

DEFAULT_MODEL = "all-MiniLM-L6-v2"


def _get_model(model_name: str = DEFAULT_MODEL) -> "SentenceTransformer":
    """Lazy-load and cache the embedding model."""
    if model_name not in _model_cache:
        from sentence_transformers import SentenceTransformer

        _model_cache[model_name] = SentenceTransformer(model_name)
    return _model_cache[model_name]


def embed_text(text: str, model_name: str = DEFAULT_MODEL) -> List[float]:
    """Generate an embedding vector for a single text string."""
    model = _get_model(model_name)
    vec = model.encode(text, convert_to_numpy=True)
    return vec.tolist()


def embed_batch(texts: List[str], model_name: str = DEFAULT_MODEL) -> List[List[float]]:
    """Generate embeddings for a batch of texts."""
    model = _get_model(model_name)
    vecs = model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
    return [v.tolist() for v in vecs]


class Indexer:
    """Manages embedding generation and delegates search to the store."""

    def __init__(self, model_name: str = DEFAULT_MODEL) -> None:
        self.model_name = model_name

    def embed(self, text: str) -> List[float]:
        return embed_text(text, self.model_name)

    def embed_many(self, texts: List[str]) -> List[List[float]]:
        return embed_batch(texts, self.model_name)
