"""FastEmbed-based embedding provider (CPU-friendly)."""

from __future__ import annotations

from context_manager.prune.base import Embedder

# Default model: small, fast, and good quality for deduplication.
_DEFAULT_MODEL = "BAAI/bge-small-en-v1.5"


class FastEmbedEmbedder(Embedder):
    """Generate embeddings using the ``fastembed`` library.

    ``fastembed`` is a lightweight, CPU-friendly embedding library that
    requires no heavy PyTorch installation.

    Args:
        model_name: The fastembed model identifier.
            Defaults to ``"BAAI/bge-small-en-v1.5"``.
    """

    def __init__(self, model_name: str = _DEFAULT_MODEL) -> None:
        self._model_name = model_name
        self._model = None  # lazy init

    def _get_model(self):
        """Lazily initialise the fastembed ``TextEmbedding`` model."""
        if self._model is None:
            try:
                from fastembed import TextEmbedding
            except ImportError as exc:
                raise ImportError(
                    "fastembed is required for FastEmbedEmbedder. "
                    "Install it with: pip install llm-ctx-mgr[prune]"
                ) from exc
            self._model = TextEmbedding(model_name=self._model_name)
        return self._model

    def embed(self, texts: list[str]) -> list[list[float]]:
        """Return embedding vectors for each string in *texts*.

        Returns:
            A list of float lists, one per input text.
        """
        if not texts:
            return []
        model = self._get_model()
        # fastembed returns a generator of numpy arrays.
        return [vector.tolist() for vector in model.embed(texts)]

    def model_name(self) -> str:
        """Return the fastembed model identifier."""
        return self._model_name
