import os
from langchain_community.embeddings import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(
        collection_name="knowledge",
        embedding_function=embeddings,
        persist_directory=persist_directory,
    )

def load_documents(directory: str, vectorstore: Chroma) -> None:
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    documents = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith((".txt", ".md")):
                path = os.path.join(root, file)
                with open(path, "r", encoding="utf-8") as f:
                    content = f.read()
                docs = splitter.split_text(content)
                for i, chunk in enumerate(docs):
                    documents.append(
                        Document(
                            page_content=chunk,
                            metadata={"source": path, "chunk": i},
                        )
                    )
    if documents:
        vectorstore.add_documents(documents)
        vectorstore.persist()