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

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

    def _load_model(self):
        print(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):
                    print("    Corrupted cache detected. Attempting repair...")
                    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 source: {self.source}")

    def get_function(self): return self.model




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

# class EmbeddingsManager:
#     def __init__(self, 
#                  model_name="BAAI/bge-small-en-v1.5", 
#                  source="fastembed", 
#                  device="cpu", 
#                  api_key=None):
        
#         self.source = source.lower()
#         self.model_name = model_name
#         self.device = device
#         self.api_key = api_key
        
#         # Désactivation explicite de hf_transfer sur Windows pour éviter la corruption de fichiers
#         # Le téléchargeur standard est plus lent mais beaucoup plus fiable.
#         os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
        
#         self.model = self._load_model()

#     def _load_model(self):
#         print(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:
#                 # Gestion automatique du cache corrompu
#                 if "NO_SUCHFILE" in str(e) or "File doesn't exist" in str(e):
#                     print("⚠️ Corrupted cache detected. Attempting repair...")
#                     # On supprime le cache et on réessaie une fois
#                     if os.path.exists(cache_path):
#                         shutil.rmtree(cache_path)
#                         os.makedirs(cache_path)
                    
#                     print("🔄 Re-downloading the model (Standard mode)...")
#                     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 source: {self.source}")

#     def get_function(self):
#         return self.model