"""Write an annotated PDF: the anatomy-snap table regions drawn on each page.

Runs the recommended approach (vision+grid anatomy prompt, Haiku@temp0, with
text-layer title-snap) over every page and draws each region as a rectangle
on the original page, plus a small ordinal label, then saves a new PDF so the
boundaries can be reviewed in a normal PDF viewer.
"""

from __future__ import annotations

import argparse
import asyncio
import tempfile
from pathlib import Path

import fitz

from experiments.que245.approaches import make_locator

OUT = Path("experiments/que245/render")


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pdfs", nargs="+")
    ap.add_argument("--rows", type=int, default=36)
    ap.add_argument("--cols", type=int, default=12)
    ap.add_argument("--approach", default="anatomy-snap")
    ap.add_argument("--model", default="claude-haiku-4-5-20251001")
    args = ap.parse_args()
    OUT.mkdir(parents=True, exist_ok=True)

    loc = make_locator(args.approach, model=args.model, temperature=0.0, rows=args.rows, cols=args.cols)
    blue = (0.0, 0.32, 0.92)

    for pdf_path in args.pdfs:
        pdf = Path(pdf_path)
        doc = fitz.open(str(pdf))
        total = 0
        for pno in range(doc.page_count):
            page = doc[pno]
            w, h = page.rect.width, page.rect.height
            with tempfile.TemporaryDirectory() as tmp:
                ip = Path(tmp) / "p.png"
                page.get_pixmap(dpi=200).save(str(ip))
                located = await loc.locate(ip, pdf, pno + 1)
            for t in located:
                x0, y0, x1, y1 = t.region
                rect = fitz.Rect(min(x0, x1) * w, min(y0, y1) * h, max(x0, x1) * w, max(y0, y1) * h)
                page.draw_rect(rect, color=blue, width=1.4)
                label = f"{t.ordinal}: {t.title}" if t.title else str(t.ordinal)
                page.insert_text((rect.x0 + 2, max(8, rect.y0 - 3)), label, fontsize=6, color=blue)
            total += len(located)
            print(f"  {pdf.name} p{pno+1}: {len(located)} tables")
        out = OUT / f"boxed_{pdf.stem[:28]}.pdf"
        doc.save(str(out))
        print(f"saved {out}  ({total} boxes over {doc.page_count} pages)\n")


if __name__ == "__main__":
    asyncio.run(main())