import os
from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader, DirectoryLoader


def create_vectorstore(persist_directory: str = "./chroma_db"):
    """Create a ChromaDB vectorstore with Ollama embeddings."""
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embeddings
    )
    return vectorstore


def load_documents(directory: str, vectorstore: Chroma):
    """Load documents from directory, chunk them, and add to vectorstore."""
    # Load all .txt and .md files from directory
    loader = DirectoryLoader(
        directory,
        glob="**/*.{txt,md}",
        loader_cls=TextLoader,
        loader_kwargs={"encoding": "utf-8"}
    )
    documents = loader.load()
    
    # Chunk documents
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
        is_separator_regex=False
    )
    chunks = text_splitter.split_documents(documents)
    
    # Add to vectorstore
    if chunks:
        vectorstore.add_documents(chunks)
        print(f"Loaded {len(chunks)} chunks from {len(documents)} documents")
    else:
        print("No documents found to load")
    
    return vectorstore