import os
import time
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader

folder_path = "documents"

def load_pdfs(folder_path):
    docs = []
    for filename in os.listdir(folder_path):
        if filename.endswith(".pdf"):
            print(f"Processing: {filename}")
            loader = PyPDFLoader(os.path.join(folder_path, filename))
            pages = loader.load()

            for page in pages:
                page.metadata["source"] = filename

            docs.extend(pages)

    print(f"Total documents loaded: {len(docs)}")
    return docs


def split_documents(docs):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50
    )
    chunks = text_splitter.split_documents(docs)
    print(f"Total chunks: {len(chunks)}")
    return chunks


def create_vectorstore(docs):
    embedding = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    print("Creating vector DB...")
    start = time.time()

    vector_db = Chroma.from_documents(
        docs,
        embedding,
        persist_directory="./chroma_db"
    )

    vector_db.persist()

    print(f"Done in {time.time() - start} sec")
    return vector_db


def main():
    docs = load_pdfs(folder_path)
    chunks = split_documents(docs)
    vector_db = create_vectorstore(chunks)

    print(f"Stored documents: {vector_db._collection.count()}")


if __name__ == "__main__":
    main()