import uuid
from langchain_chroma import Chroma
from langchain_core.stores import InMemoryStore
from langchain_core.documents import Document
from langchain_classic.retrievers.multi_vector import MultiVectorRetriever
from langchain_huggingface import HuggingFaceEmbeddings

def initialize_retriever():
    vectorstore = Chroma(
        collection_name="multi_modal_rag",
        embedding_function=HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    )
    store = InMemoryStore()
    id_key = "doc_id"
    retriever = MultiVectorRetriever(vectorstore=vectorstore, docstore=store, id_key=id_key)
    return retriever, id_key

def add_assets_to_vectorstore(retriever, id_key, raw_elements, summaries):
    if not raw_elements:
        return
    doc_ids = [str(uuid.uuid4()) for _ in raw_elements]
    summary_documents = [
        Document(page_content=summary, metadata={id_key: doc_ids[i]})
        for i, summary in enumerate(summaries)
    ]
    retriever.vectorstore.add_documents(summary_documents)
    retriever.docstore.mset(list(zip(doc_ids, raw_elements)))

def run_database_pipeline(text_chunks, text_sums, tables, table_sums, images, image_sums):
    retriever, id_key = initialize_retriever()
    add_assets_to_vectorstore(retriever, id_key, text_chunks, text_sums)
    add_assets_to_vectorstore(retriever, id_key, tables, table_sums)
    add_assets_to_vectorstore(retriever, id_key, images, image_sums)
    return retriever