"""Draw both dual-flow result sets onto the source PDF.

Reads a DualResult JSON (from `quber dual`) and overlays, per page:
- VISUAL flow boxes in blue  (normalized 0..1, top-left origin) labeled V#: title
- CAMELOT flow boxes in red  (PDF points, bottom-left origin -> converted)
  labeled with the is_table verdict, so the two independent flows and their
  divergence are visible on one page.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import fitz

VISUAL = (0.0, 0.32, 0.92)
CAMELOT_TABLE = (0.85, 0.10, 0.10)
CAMELOT_NOT = (0.95, 0.55, 0.0)
OUT = Path("experiments/que246/render")


def draw(dual_json: Path, source: Path) -> Path:
    data = json.loads(dual_json.read_text())
    doc = fitz.open(str(source))

    visual_by_page: dict[int, list] = {}
    for v in data["visual"]:
        visual_by_page.setdefault(v["page"], []).append(v)
    camelot_by_page: dict[int, list] = {}
    for c in data["camelot"]:
        camelot_by_page.setdefault(c["candidate"]["page"], []).append(c)

    for pno in range(doc.page_count):
        page = doc[pno]
        w, h = page.rect.width, page.rect.height
        for v in visual_by_page.get(pno + 1, []):
            x0, y0, x1, y1 = v["region"]
            rect = fitz.Rect(x0 * w, y0 * h, x1 * w, y1 * h)
            page.draw_rect(rect, color=VISUAL, width=1.6)
            page.insert_text((rect.x0 + 2, max(8, rect.y0 - 3)),
                             f"V{v['ordinal']}: {v['title']}", fontsize=6, color=VISUAL)
        for c in camelot_by_page.get(pno + 1, []):
            bbox = c["candidate"]["bbox"]
            if not bbox:
                continue
            x1c, y1c, x2c, y2c = bbox  # PDF points, bottom-left origin
            rect = fitz.Rect(x1c, h - y2c, x2c, h - y1c)
            cls = c.get("classification")
            is_table = cls.get("is_table") if cls else None
            color = CAMELOT_TABLE if is_table else CAMELOT_NOT
            page.draw_rect(rect, color=color, width=1.0, dashes="[3] 0")
            page.insert_text((rect.x0 + 2, min(h - 4, rect.y1 + 9)),
                             f"C {c['candidate']['flavor']} is_table={is_table}", fontsize=6, color=color)

    OUT.mkdir(parents=True, exist_ok=True)
    out = OUT / f"dual_boxed_{source.stem[:24]}.pdf"
    doc.save(str(out))
    return out


def main():
    pairs = [
        ("BHE_991.dual.json", "documents/BHE_991.pdf"),
        ("Q2FY25-Visa-Operational-Performance-Data-FINAL.dual.json",
         "documents/Q2FY25-Visa-Operational-Performance-Data-FINAL.pdf"),
        ("TMUS_Q225_991.dual.json", "documents/TMUS_Q225_991.pdf"),
    ]
    tmp = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/home/mande/.claude/jobs/a95a5f59/tmp")
    for js, pdf in pairs:
        out = draw(tmp / js, Path(pdf))
        print(f"saved {out}")


if __name__ == "__main__":
    main()