import os
import shutil
import hashlib
from langchain_chroma import Chroma
from langchain_community.vectorstores import FAISS, Qdrant

try: 
    from qdrant_client import QdrantClient
    from qdrant_client.http import models as qdrant_models
except ImportError: QdrantClient = None

class VectorDBFactory:
    @staticmethod
    def reset(persist_directory="./db_storage"):
        if os.path.exists(persist_directory):
            try:
                shutil.rmtree(persist_directory)
                print(f"      DB folder deleted: {persist_directory}")
            except Exception as e: print(f"    DB cleanup error: {e}")

    @staticmethod
    def create(db_type, embedding_function, persist_directory="./db_storage"):
        db_type = db_type.lower()
        if db_type == "chroma":
            return Chroma(persist_directory=persist_directory, embedding_function=embedding_function)
        elif db_type == "faiss":
            try: return FAISS.load_local(persist_directory, embedding_function, allow_dangerous_deserialization=True)
            except: return None 
        elif db_type == "qdrant":
            if os.path.exists(persist_directory):
                try:
                    client = QdrantClient(path=persist_directory)
                    if any(c.name == "rostaing_collection" for c in client.get_collections().collections):
                        return Qdrant(client=client, collection_name="rostaing_collection", embeddings=embedding_function, content_payload_key="page_content", metadata_payload_key="metadata")
                except: return None
            return None
        else: raise ValueError(f"Database '{db_type}' not supported.")

    @staticmethod
    def _generate_ids(docs):
        ids = []
        for doc in docs:
            content_str = f"{doc.page_content}-{doc.metadata.get('source', '')}"
            hash_id = hashlib.md5(content_str.encode('utf-8')).hexdigest()
            ids.append(hash_id)
        return ids

    @staticmethod
    def add_docs(db_instance, docs, db_type, embedding_function, persist_directory="./db_storage"):
        if not docs: return db_instance
        db_type = db_type.lower()
        ids = VectorDBFactory._generate_ids(docs)

        if db_instance is None:
            print(f"    Initial database creation {db_type.upper()}...")
            if db_type == "faiss":
                db = FAISS.from_documents(docs, embedding_function)
                db.save_local(persist_directory)
                return db
            elif db_type == "qdrant":
                if not QdrantClient: raise ImportError("pip install qdrant-client")
                client = QdrantClient(path=persist_directory)
                client.recreate_collection(collection_name="rostaing_collection", vectors_config=qdrant_models.VectorParams(size=len(embedding_function.embed_query("test")), distance=qdrant_models.Distance.COSINE))
                db = Qdrant(client=client, collection_name="rostaing_collection", embeddings=embedding_function, content_payload_key="page_content", metadata_payload_key="metadata")
                db.add_documents(docs, ids=ids)
                return db
            elif db_type == "chroma":
                return Chroma.from_documents(docs, embedding_function, persist_directory=persist_directory, ids=ids)
        else:
            try:
                if db_type == "faiss": 
                    db_instance.add_documents(docs)
                    db_instance.save_local(persist_directory)
                else: db_instance.add_documents(docs, ids=ids)
                return db_instance
            except Exception as e:
                print(f"  Error while adding to the DB: {e}")
                return db_instance

# # -----------------------------------------------------
# import os
# import shutil
# import hashlib

# # Imports LangChain pour les bases locales
# from langchain_chroma import Chroma
# from langchain_community.vectorstores import FAISS, Qdrant

# # Import Client Qdrant
# try:
#     from qdrant_client import QdrantClient
#     from qdrant_client.http import models as qdrant_models
# except ImportError:
#     QdrantClient = None

# class VectorDBFactory:
#     @staticmethod
#     def reset(persist_directory="./db_storage"):
#         """Physical cleanup of the storage folder."""
#         if os.path.exists(persist_directory):
#             try:
#                 def on_rm_error(func, path, exc_info):
#                     try:
#                         os.chmod(path, 0o777)
#                         os.remove(path)
#                     except: pass
#                 shutil.rmtree(persist_directory, onerror=on_rm_error)
#                 print(f"🧹 DB folder deleted: {persist_directory}")
#             except Exception as e:
#                 print(f"⚠️ DB cleanup error: {e}")

#     @staticmethod
#     def create(db_type, embedding_function, persist_directory="./db_storage"):
#         """Load an existing DB."""
#         db_type = db_type.lower()
        
#         if db_type == "chroma":
#             return Chroma(persist_directory=persist_directory, embedding_function=embedding_function)
        
#         elif db_type == "faiss":
#             try:
#                 return FAISS.load_local(persist_directory, embedding_function, allow_dangerous_deserialization=True)
#             except: return None 

#         elif db_type == "qdrant":
#             if os.path.exists(persist_directory):
#                 try:
#                     client = QdrantClient(path=persist_directory)
#                     cols = client.get_collections().collections
#                     if any(c.name == "rostaing_collection" for c in cols):
#                         return Qdrant(
#                             client=client, 
#                             collection_name="rostaing_collection", 
#                             embeddings=embedding_function,
#                             content_payload_key="page_content",
#                             metadata_payload_key="metadata"
#                         )
#                 except: return None
#             return None

#         else:
#             raise ValueError(f"Database '{db_type}' not supported.")

#     @staticmethod
#     def _generate_ids(docs):
#         """Generate unique IDs based on content to avoid duplicates."""
#         ids = []
#         for doc in docs:
#             # Hash du contenu + source pour l'unicité
#             content_str = f"{doc.page_content}-{doc.metadata.get('source', '')}"
#             hash_id = hashlib.md5(content_str.encode('utf-8')).hexdigest()
#             ids.append(hash_id)
#         return ids

#     @staticmethod
#     def add_docs(db_instance, docs, db_type, embedding_function, persist_directory="./db_storage"):
#         """Add documents with deduplication (IDs)."""
#         if not docs: return db_instance
#         db_type = db_type.lower()
        
#         # Génération des IDs uniques pour éviter la duplication si reset_db=False
#         ids = VectorDBFactory._generate_ids(docs)

#         # --- CAS 1 : CRÉATION INITIALE ---
#         if db_instance is None:
#             print(f"✨ Initial database creation {db_type.upper()}...")
            
#             if db_type == "faiss":
#                 # FAISS ne gère pas nativement les IDs aussi facilement pour la dédup, 
#                 # mais on initialise propre.
#                 db = FAISS.from_documents(docs, embedding_function)
#                 db.save_local(persist_directory)
#                 return db

#             elif db_type == "qdrant":
#                 if not QdrantClient: raise ImportError("pip install qdrant-client")
                
#                 dummy_vec = embedding_function.embed_query("test")
#                 vector_size = len(dummy_vec)
                
#                 client = QdrantClient(path=persist_directory)
#                 client.recreate_collection(
#                     collection_name="rostaing_collection",
#                     vectors_config=qdrant_models.VectorParams(
#                         size=vector_size, 
#                         distance=qdrant_models.Distance.COSINE
#                     )
#                 )
                
#                 db = Qdrant(
#                     client=client, 
#                     collection_name="rostaing_collection", 
#                     embeddings=embedding_function,
#                     content_payload_key="page_content",
#                     metadata_payload_key="metadata"
#                 )
#                 # Ajout avec IDs
#                 db.add_documents(docs, ids=ids)
#                 return db
            
#             elif db_type == "chroma":
#                 # Chroma gère très bien les IDs pour l'upsert (mise à jour si existe)
#                 return Chroma.from_documents(
#                     docs, 
#                     embedding_function, 
#                     persist_directory=persist_directory,
#                     ids=ids
#                 )

#         # --- CAS 2 : MISE À JOUR ---
#         else:
#             try:
#                 # L'ajout avec IDs permet à Chroma/Qdrant d'ignorer ou mettre à jour les doublons
#                 # au lieu de les ajouter bêtement.
#                 if db_type == "faiss":
#                     # Pour FAISS, on ajoute tout (la dédup est plus complexe sur fichier plat)
#                     db_instance.add_documents(docs)
#                     db_instance.save_local(persist_directory)
#                 else:
#                     db_instance.add_documents(docs, ids=ids)
                    
#                 return db_instance
#             except Exception as e:
#                 print(f"❌ Error while adding to the DB: {e}")
#                 return db_instance