#!/usr/bin/env python3
"""
OCR all PDFs that have associated .py recipe files in convergence/data/recipes/generated.
Saves extracted text to temp/ocr-fuzzy/output/<category>__<id>.txt

Uses multiprocessing to parallelize across PDFs.
"""

import multiprocessing as mp
import sys
import time
from pathlib import Path

RECIPES_DIR = Path("/Users/soma/Development/convergence/data/recipes/generated")
PDFS_DIR = Path("/Users/soma/Development/convergence/data/pdfs")
OUTPUT_DIR = Path(__file__).parent / "output"


def discover_pdfs():
    """Find all category/id pairs that have at least one .py recipe file."""
    seen = set()
    results = []
    for py_file in RECIPES_DIR.rglob("*.py"):
        # Structure: generated/<category>/<id>/<recipe>.py
        id_dir = py_file.parent
        category_dir = id_dir.parent
        if category_dir.parent != RECIPES_DIR:
            continue  # skip nested debug dirs etc.
        key = (category_dir.name, id_dir.name)
        if key in seen:
            continue
        seen.add(key)
        pdf_path = PDFS_DIR / category_dir.name / f"{id_dir.name}.pdf"
        if pdf_path.exists():
            results.append((category_dir.name, id_dir.name, pdf_path))
        else:
            print(f"WARNING: PDF not found: {pdf_path}")
    return sorted(results)


def process_pdf(args):
    category, doc_id, pdf_path = args
    output_file = OUTPUT_DIR / f"{category}__{doc_id}.txt"
    if output_file.exists():
        return f"SKIP {category}/{doc_id} (already done)"

    try:
        from natural_pdf import PDF

        pdf = PDF(str(pdf_path), text_layer=False)
        try:
            pdf.apply_ocr()
            text = pdf.extract_text()
            output_file.write_text(text, encoding="utf-8")
            return f"OK   {category}/{doc_id} ({len(pdf.pages)} pages)"
        finally:
            pdf.close()
    except Exception as e:
        return f"FAIL {category}/{doc_id}: {e}"


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    pdfs = discover_pdfs()
    print(f"Found {len(pdfs)} PDFs to process")

    workers = min(mp.cpu_count(), 8)
    print(f"Using {workers} workers\n")

    t0 = time.time()
    with mp.Pool(workers) as pool:
        for result in pool.imap_unordered(process_pdf, pdfs):
            print(result)

    elapsed = time.time() - t0
    print(f"\nDone in {elapsed:.1f}s")


if __name__ == "__main__":
    main()
