#!/usr/bin/env python3
"""Add all PDFs under a folder to the RAG vector store. Run from project root.

The settings below are the defaults; they can be overridden on the command line, e.g.
    python add_pdfs.py --data data --index ./vectorstore_docling --loader docling
"""
import argparse
import shutil
from pathlib import Path

from local_rag import CHUNK_OVERLAP, CHUNK_SIZE, PDF_LOADER, PROCESSED_DIR, LocalRAG

# Folder to scan for PDFs (recursively).
DATA_ROOT = Path("/Users/Philipp/Desktop/workspace/python/gpt_publikationen/data_vs")

# The chat uses the vector stores named in personas.py (currently ./vectorstore_pypdf_2000); pass --index to build one.
VECTORSTORE_PATH = "./vectorstore"

# If True: delete the existing vector store folder, then index all PDFs from scratch.
# If False: load the existing index (if any) and merge new chunks from these PDFs into it.
CLEAR_VECTORSTORE_FIRST = True


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--data", type=Path, default=DATA_ROOT, help="folder to scan for PDFs")
    parser.add_argument("--index", default=VECTORSTORE_PATH, help="vector store folder")
    parser.add_argument("--loader", default=PDF_LOADER, choices=["pypdf", "docling"])
    parser.add_argument("--processed", default=PROCESSED_DIR, help="Docling output (prepare_corpus.py)")
    parser.add_argument("--chunk-size", type=int, default=CHUNK_SIZE)
    parser.add_argument("--chunk-overlap", type=int, default=CHUNK_OVERLAP)
    parser.add_argument("--no-headers", action="store_true", help="no author/title header per chunk")
    args = parser.parse_args()
    DATA_ROOT, VECTORSTORE_PATH = args.data, args.index

    if CLEAR_VECTORSTORE_FIRST and Path(VECTORSTORE_PATH).exists():
        shutil.rmtree(VECTORSTORE_PATH)
        print(f"Cleared existing vector store: {VECTORSTORE_PATH}")
    elif not CLEAR_VECTORSTORE_FIRST:
        print(f"Appending to existing vector store (if any): {VECTORSTORE_PATH}")

    pdfs = sorted(p for p in DATA_ROOT.rglob("*") if p.suffix.lower() == ".pdf")
    print(f"Found {len(pdfs)} PDF(s) under {DATA_ROOT}")
    if not pdfs:
        raise SystemExit("No PDFs found.")

    rag = LocalRAG(vectorstore_path=VECTORSTORE_PATH, pdf_loader=args.loader,
                   processed_dir=args.processed, chunk_size=args.chunk_size,
                   chunk_overlap=args.chunk_overlap, chunk_headers=not args.no_headers)
    rag.add_documents([str(p) for p in pdfs])
    print("Done.")