from langchain_community.document_loaders import PyPDFLoader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter


def ingest_documents(pdf_path: str):
    loader = PyPDFLoader(pdf_path)
    documents = loader.load()

    print(f"{len(documents)} Dokumente eingelesen")

    return documents


def generate_chunks(documents):
    splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
    chunks = splitter.split_documents(documents)

    return chunks


def generate_embeddings(chunks, persist_dir: str = "./chroma_db"):
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    vectorstore = Chroma.from_documents(
        documents=chunks, embedding=embeddings, persist_directory=persist_dir
    )

    print(f"{len(chunks)} Chunks gespeichert")

    return vectorstore


if __name__ == "__main__":
    documents = ingest_documents("docs/Stefan_Schneider_Lebenslauf.pdf")
    chunks = generate_chunks(documents)
    store = generate_embeddings(chunks)