import os
from pathlib import Path

# Important: make sure dotenv is loaded so Cohere API key is available
from dotenv import load_dotenv
load_dotenv()

# Import the vector store function and path from main.py
from main import get_vector_store, FAISS_DIR

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader

def ingest_folder(folder_path):
    vs = get_vector_store()
    folder = Path(folder_path)
    
    if not folder.exists():
        print(f"Folder '{folder_path}' does not exist.")
        return

    docs_to_add = []
    
    print(f"Scanning '{folder_path}' for documents...")
    
    for file_path in folder.rglob("*"):
        if not file_path.is_file():
            continue
            
        ext = file_path.suffix.lower()
        try:
            if ext == ".pdf":
                loader = PyPDFLoader(str(file_path))
            elif ext in (".txt", ".md"):
                loader = TextLoader(str(file_path), encoding="utf-8")
            elif ext == ".docx":
                loader = Docx2txtLoader(str(file_path))
            else:
                continue
                
            print(f"Loading {file_path.name}...")
            docs = loader.load()
            for d in docs:
                d.metadata["source"] = file_path.name
            
            splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
            chunks = splitter.split_documents(docs)
            docs_to_add.extend(chunks)
        except Exception as e:
            print(f"Error loading {file_path.name}: {e}")
            
    if docs_to_add:
        print(f"Adding {len(docs_to_add)} chunks to FAISS index...")
        vs.add_documents(docs_to_add)
        vs.save_local(str(FAISS_DIR))
        print("Done! Documents successfully indexed.")
    else:
        print("No new supported documents found to add.")

if __name__ == "__main__":
    ingest_folder("docs")