import os
from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter

class VectorStore:
    def __init__(self, persist_path: str = "vectorstore"):
        self.persist_path = persist_path
        os.makedirs(self.persist_path, exist_ok=True)
        embeddings = OllamaEmbeddings(model="nomic-embed-text")
        self.store = Chroma(persist_directory=self.persist_path, embedding_function=embeddings)

    def add_documents(self, docs_dir: str):
        texts = []
        for root, _, files in os.walk(docs_dir):
            for fname in files:
                if fname.lower().endswith((".txt", ".md")):
                    path = os.path.join(root, fname)
                    with open(path, "r", encoding="utf-8") as f:
                        content = f.read()
                    texts.append(content)
        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        chunks = splitter.split_text("\n".join(texts))
        self.store.add_texts(chunks)
        self.store.persist()

    def get_store(self):
        return self.store