import os
from typing import Union, List
import numpy as np
from sentence_transformers import SentenceTransformer
from .base import BaseEncoder


class TextEncoder(BaseEncoder):
    """Text encoder using SentenceTransformers."""

    def __init__(self, model_name: str = None):
        self.model_name = model_name or os.getenv(
            "HIEROMEM_ENCODER_MODEL", "all-MiniLM-L6-v2")
        self._model = SentenceTransformer(self.model_name)

    def encode(self, data: Union[str, List[str]]) -> np.ndarray:
        """
        Encode text into a vector.

        Args:
            data: Text string or list of strings

        Returns:
            Numpy array representing the embedding vector
        """
        return self._model.encode(data, convert_to_numpy=True)

    @property
    def dimension(self) -> int:
        """Return the dimension of the embedding vector."""
        return self._model.get_sentence_embedding_dimension()
