# ------------------------------------------------------------
# prepare_pdf.py
# ------------------------------------------------------------
# Utility script for running PDF→Parquet conversions
# either for a predefined list of documents or for a single file.
#
# Mode 1 (recommended):
#   Fill the CHAPTERS/DOCUMENTS list and run:
#       python prepare_pdf.py
#
# Mode 2 (single PDF from CLI):
#       python prepare_pdf.py --pdf "/path/to/notebook.pdf" --doc-id "DOC-001-File-Notebook"
#
# Output: Each run produces Parquet embeddings and a JSONL manifest
#
# Requirements:
#   Set the same API key environment variable as used in the project:
#       export EMBEDDING_API_KEY=...
# ------------------------------------------------------------

import argparse
from typing import List, Tuple
from build_embeddings import build_for_pdf

# ====== List of documents with (doc_id, path) ======
CHAPTERS: List[Tuple[str, str]] = [
    # Example entries:
    # ("DOC-001-Employee-Handbook", r"/path/to/Employee-Handbook.pdf"),
    # ("DOC-002-Product-Manual", r"/path/to/Product-Manual.pdf"),
    # ("DOC-003-Research-Paper", r"/path/to/Research-Paper.pdf"),

    # Leave empty if running with --pdf and --doc-id arguments only.
]

def run_single(pdf: str, doc_id: str) -> None:
    print(f"[RUN] {doc_id} ← {pdf}")
    out = build_for_pdf(pdf, doc_id)
    print(f"[OK ] {doc_id} → {out}")

def run_batch(pairs: List[Tuple[str, str]]) -> None:
    if not pairs:
        print("[WARN] CHAPTERS list is empty. Fill it or use --pdf/--doc-id parameters.")
        return
    errors = 0
    for doc_id, pdf in pairs:
        try:
            run_single(pdf, doc_id)
        except Exception as e:
            errors += 1
            print(f"[FAIL] {doc_id}: {e}")
    print(f"\n[DONE] Run complete. Total: {len(pairs)}, Success: {len(pairs) - errors}, Failures: {errors}")

def main():
    parser = argparse.ArgumentParser(description="Prepare PDFs → Parquet (single or batch).")
    parser.add_argument("--pdf", help="Path to a single PDF file.")
    parser.add_argument("--doc-id", help="Document identifier, e.g. DOC-001-Employee-Handbook.")
    args = parser.parse_args()

    if args.pdf and args.doc_id:
        run_single(args.pdf, args.doc_id)
    elif args.pdf or args.doc_id:
        parser.error("Both --pdf and --doc-id must be provided together, or neither (use CHAPTERS list).")
    else:
        run_batch(CHAPTERS)

if __name__ == "__main__":
    main()