import os
from pathlib import Path
from langchain.vectorstores import Chroma
from langchain.embeddings.ollama import OllamaEmbeddings
from langchain.document_loaders import TextLoader, MarkdownLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

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

def load_documents(directory: str, vectorstore):
    docs = []
    for file in Path(directory).glob("**/*.*"):
        if file.suffix.lower() not in {".txt", ".md"}:
            continue
        loader = TextLoader(str(file)) if file.suffix == ".txt" else MarkdownLoader(str(file))
        docs.extend(loader.load())
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    texts = splitter.split_documents(docs)
    vectorstore.add_documents(texts)
    vectorstore.persist()