import os
import shutil
from langchain_community.embeddings import FastEmbedEmbeddings
from langchain_ollama import OllamaEmbeddings
from langchain_huggingface import HuggingFaceEmbeddings

class EmbeddingsManager:
    """Manages Vector Embeddings efficiently."""
    def __init__(self, model_name="BAAI/bge-small-en-v1.5", source="fastembed", device="cpu", api_key=None, logger=None):
        self.source = source.lower()
        self.model_name = model_name
        self.device = device
        self.api_key = api_key
        self.logger = logger
        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
        self.model = self._load_model()

    def _log(self, msg, level="info"):
        if self.logger: self.logger(msg, level)
        else: print(msg)

    def _load_model(self):
        self._log(f"Loading embeddings: {self.source} ({self.model_name})...")
        if self.source == "fastembed":
            cache_path = os.path.abspath("./embedding_cache")
            if not os.path.exists(cache_path): os.makedirs(cache_path)
            try:
                return FastEmbedEmbeddings(model_name=self.model_name, threads=None, cache_dir=cache_path)
            except Exception as e:
                if "NO_SUCHFILE" in str(e) or "File doesn't exist" in str(e):
                    self._log("Corrupted cache detected. Attempting repair...", "warning")
                    if os.path.exists(cache_path): shutil.rmtree(cache_path); os.makedirs(cache_path)
                    return FastEmbedEmbeddings(model_name=self.model_name, threads=None, cache_dir=cache_path)
                else: raise e
        elif self.source == "ollama":
            return OllamaEmbeddings(model=self.model_name, base_url="http://localhost:11434")
        elif self.source == "huggingface":
            return HuggingFaceEmbeddings(model_name=self.model_name, model_kwargs={'device': self.device}, encode_kwargs={'normalize_embeddings': True})
        elif self.source == "openai":
            try: 
                from langchain_openai import OpenAIEmbeddings
                return OpenAIEmbeddings(model=self.model_name or "text-embedding-3-small", api_key=self.api_key or os.getenv("OPENAI_API_KEY"))
            except ImportError: raise ImportError("Install: pip install langchain-openai")
        else: raise ValueError(f"Unknown embedding source: {self.source}")

    def get_function(self): return self.model