import os
from rag_store import KnowledgeBase
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

def load_documents_from_dir(directory: str, kb: KnowledgeBase):
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(('.txt', '.md')):
                path = os.path.join(root, file)
                try:
                    with open(path, 'r', encoding='utf-8') as f:
                        text = f.read()
                    docs = [Document(page_content=text, metadata={"title": file})]
                    chunks = splitter.split_documents(docs)
                    kb.add_documents(chunks)
                    print(f"Loaded {file} into KB.")
                except Exception as e:
                    print(f"Failed to load {file}: {e}")

if __name__ == "__main__":
    kb = KnowledgeBase()
    load_documents_from_dir("./data", kb)