from __future__ import annotations

from pathlib import Path
from typing import List, Optional, TYPE_CHECKING
import os

from analogai.foundation.common.logging.app_logger import app_logger

if TYPE_CHECKING:
    from sentence_transformers import SentenceTransformer


class LocalSentenceTransformerProvider:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2", cache_folder: Optional[str] = None):
        self.model_name = model_name
        self._model: Optional["SentenceTransformer"] = None
        self._dimension: Optional[int] = None
        self.cache_folder = cache_folder or os.getenv(
            "SENTENCE_TRANSFORMERS_CACHE",
            str(Path.home() / ".cache" / "sentence_transformers"),
        )
        Path(self.cache_folder).mkdir(parents=True, exist_ok=True)
        app_logger.info(
            f"Local embeddings initializing with model: {model_name}, cache: {self.cache_folder}"
        )
        from sentence_transformers import SentenceTransformer

        app_logger.info(f"Loading model {self.model_name} from cache: {self.cache_folder}")
        self._model = SentenceTransformer(self.model_name, cache_folder=self.cache_folder)
        self._dimension = self._model.get_sentence_embedding_dimension()
        app_logger.info(f"Local embedding model loaded with dimension: {self._dimension}")

    @property
    def model(self) -> "SentenceTransformer":
        if self._model is None:
            raise RuntimeError(
                "Local embedding model not initialized. "
                "Instantiate only when USE_VECTOR_SEARCH is enabled."
            )
        return self._model

    @property
    def dimension(self) -> int | None:
        if self._dimension is None:
            if self._model is None:
                return None
            self._dimension = self._model.get_sentence_embedding_dimension()
        return self._dimension

    def encode(self, text: str) -> List[float]:
        if not text or not text.strip():
            if self.dimension is None:
                return []
            app_logger.warning("Attempted to encode empty text, returning zero vector")
            return [0.0] * self.dimension
        return self.model.encode(text, convert_to_numpy=True).tolist()

    def encode_batch(self, texts: List[str]) -> List[List[float]]:
        if not texts:
            return []
        filtered_texts = [text if text and text.strip() else "" for text in texts]
        return self.model.encode(filtered_texts, convert_to_numpy=True).tolist()
