#!/usr/bin/env python3
"""
rapid_pass.py — generalise the proven AQA "RapidOCR margin-pass" (95.2% on the image-only
8463 paper) to any AQA paper. Born-digital AQA QPs ship a text layer, so we force RapidOCR
over the *rendered* page (`force_ocr:true`) to simulate the image-only redistribution case
and recover the boxed `NN.M` question numbers Tesseract shatters.

For each page it writes results/<outdir>/p{N}.json (a full per-page DoclingDocument, the
shape extract.py's aqa_questions_rapid expects) and a merged.json (for board / front-matter
detection). All GPU work is serialised + OOM-resilient through dsync.

Usage:
  python scripts/rapid_pass.py samples/extra/aqa-alevel-physics-7408-1-jun22-qp.pdf rapid_7408
  python scripts/rapid_pass.py <pdf> <outdir-slug> [first_page] [last_page]
"""
import os, sys, json, subprocess, re
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import dsync

OPTS = {"ocr_engine": "rapidocr", "force_ocr": True}


def npages(pdf):
    out = subprocess.check_output(["pdfinfo", pdf]).decode()
    return int(out.split("Pages:")[1].split()[0])


def main():
    pdf = sys.argv[1]
    slug = sys.argv[2]
    if os.path.isabs(slug) or ".." in slug.split(os.sep) or not re.fullmatch(r"[A-Za-z0-9._/-]+", slug):
        raise SystemExit(f"unsafe output slug: {slug!r}")
    n = npages(pdf)
    first = int(sys.argv[3]) if len(sys.argv) > 3 else 1
    last = min(int(sys.argv[4]), n) if len(sys.argv) > 4 else n
    if first > n or first > last:
        print(f"requested page range {first}-{last} is outside PDF ({n} pages); nothing to do")
        return
    outdir = os.path.join("results", slug)
    os.makedirs(outdir, exist_ok=True)

    r = dsync._redis()
    print(f"redis: {'connected' if r else 'NO CACHE'}  pdf={pdf}  pages {first}-{last}/{n}")
    merged = {"texts": [], "tables": [], "pictures": [], "pages": {}, "_failed_pages": []}
    for pg in range(first, last + 1):
        page_path = os.path.join(outdir, f"p{pg}.json")
        if os.path.exists(page_path):
            doc = json.load(open(page_path))
            print(f"  p{pg}: file cache HIT ({len(doc.get(texts, []))} texts)")
        else:
            doc = dsync.convert_page(pdf, pg, OPTS, r=r)
        if not doc:
            merged["_failed_pages"].append(pg)
            print(f"  p{pg}: FAILED")
            continue
        json.dump(doc, open(page_path, "w"))
        for k in ("texts", "tables", "pictures"):
            merged[k].extend(doc.get(k, []))
        merged["pages"].update(doc.get("pages", {}))
        nmarg = sum(1 for t in doc.get("texts", [])
                    if (t.get("prov") or [{}])[0].get("bbox", {}).get("l", 999) <= 140)
        print(f"  p{pg}: {len(doc.get('texts', []))} texts ({nmarg} left-margin)")
    json.dump(merged, open(os.path.join(outdir, "merged.json"), "w"))
    print(f"-> {outdir}/  ({last-first+1-len(merged['_failed_pages'])} pages, "
          f"failed={merged['_failed_pages']})")


if __name__ == "__main__":
    main()
