import os
from pathlib import Path
from typing import List

from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document

def create_vectorstore(persist_directory: str = "./chroma_db"):
    if not os.path.exists(persist_directory):
        os.makedirs(persist_directory, exist_ok=True)
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)

def load_documents(directory: str, vectorstore) -> None:
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs: List[Document] = []
    for path in Path(directory).rglob("*.txt"):
        text = path.read_text(encoding="utf-8")
        docs.extend(splitter.split_documents([Document(page_content=text)]))
    for path in Path(directory).rglob("*.md"):
        text = path.read_text(encoding="utf-8")
        docs.extend(splitter.split_documents([Document(page_content=text)]))
    if docs:
        vectorstore.add_documents(docs)