"""Embedding model wrapper for Mnemosynth.

Lazy-loads sentence-transformers model on first use to keep startup fast.
Default: all-MiniLM-L6-v2 (384 dims, ~23M params, fast).
"""

from __future__ import annotations

from typing import Optional


class EmbeddingModel:
    """Wrapper around sentence-transformers for embedding generation."""

    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model_name = model_name
        self._model = None

    @property
    def model(self):
        """Lazy-load the embedding model on first use."""
        if self._model is None:
            try:
                from sentence_transformers import SentenceTransformer
                self._model = SentenceTransformer(self.model_name)
            except ImportError:
                raise ImportError(
                    "sentence-transformers is required for embeddings. "
                    "Install with: pip install sentence-transformers"
                )
        return self._model

    def embed(self, text: str) -> list[float]:
        """Generate embedding for a single text string."""
        embedding = self.model.encode(text, convert_to_numpy=True)
        return embedding.tolist()

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings for multiple texts (batched for efficiency)."""
        embeddings = self.model.encode(texts, convert_to_numpy=True, batch_size=32)
        return [e.tolist() for e in embeddings]

    @property
    def dimension(self) -> int:
        """Get the embedding dimension."""
        return self.model.get_sentence_embedding_dimension()

    def similarity(self, text1: str, text2: str) -> float:
        """Compute cosine similarity between two texts."""
        from sentence_transformers.util import cos_sim
        e1 = self.model.encode(text1, convert_to_numpy=True)
        e2 = self.model.encode(text2, convert_to_numpy=True)
        return float(cos_sim(e1, e2)[0][0])
