#!/usr/bin/env python3
"""Draw bboxes from a translated JSON onto the original PDF for visual inspection.

Usage:
    python test/draw_bbox.py --pdf input.pdf --json output.translated.json --output bbox.pdf
    python test/draw_bbox.py --pdf input.pdf --json output.json --output bbox.pdf --pages 0-9
"""

import argparse
import json
from pathlib import Path

import fitz

COLORS = {
    "FLOWING_TEXT": (0.0, 0.0, 1.0),  # blue
    "IN_PLACE": (0.0, 0.7, 0.0),  # green
    "EQUATION": (1.0, 0.5, 0.0),  # orange
    "TABLE": (0.5, 0.0, 0.8),  # purple
    "BYPASS": (1.0, 0.0, 0.0),  # red
}
CELL_COLOR = (0.8, 0.7, 0.0)  # yellow
EQ_LINE_COLOR = (0.0, 0.8, 0.8)  # cyan — equation_text_lines


def parse_pages(s: str) -> list[int] | None:
    if not s:
        return None
    result = []
    for part in s.split(","):
        part = part.strip()
        if "-" in part:
            a, b = part.split("-", 1)
            result.extend(range(int(a), int(b) + 1))
        else:
            result.append(int(part))
    return result


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--pdf", required=True)
    p.add_argument("--json", required=True)
    p.add_argument("--output", required=True)
    p.add_argument("--pages", default=None, help="e.g. 0-9 or 0,1,5")
    args = p.parse_args()

    pages_filter = parse_pages(args.pages)

    with open(args.json, encoding="utf-8") as f:
        data = json.load(f)

    doc = fitz.open(args.pdf)

    for page_idx, page_data in enumerate(data.get("pages", [])):
        if pages_filter is not None and page_idx not in pages_filter:
            continue
        if page_idx >= doc.page_count:
            break

        page = doc[page_idx]

        for elem in page_data.get("elements", []):
            cat = elem.get("category", "")
            label = elem.get("label", "")
            bbox = elem.get("bbox_pdf")
            if not bbox:
                continue

            color = COLORS.get(cat, (0.5, 0.5, 0.5))
            rect = fitz.Rect(*bbox)
            page.draw_rect(rect, color=color, width=1.5)
            page.insert_text(
                fitz.Point(bbox[0], bbox[1] - 1),
                f"{cat[:2]} {label}",
                fontsize=6,
                color=color,
            )

            if cat == "TABLE":
                for cell in elem.get("cells", []):
                    cb = cell.get("bbox_pdf")
                    if cb:
                        page.draw_rect(fitz.Rect(*cb), color=CELL_COLOR, width=0.8)

            if cat == "EQUATION":
                for line in elem.get("equation_text_lines", []) or []:
                    lb = line.get("bbox_pdf")
                    if not lb:
                        continue
                    page.draw_rect(fitz.Rect(*lb), color=EQ_LINE_COLOR, width=0.8)
                    label_text = line.get("text", "")[:20]
                    page.insert_text(
                        fitz.Point(lb[0], lb[1] - 1),
                        label_text,
                        fontsize=5,
                        color=EQ_LINE_COLOR,
                    )

    out = Path(args.output)
    out.parent.mkdir(parents=True, exist_ok=True)
    doc.save(str(out))
    doc.close()
    print(f"Saved → {out}")
    print(
        "Legend: blue=FLOWING_TEXT  green=IN_PLACE  orange=EQUATION  purple=TABLE  "
        "red=BYPASS  yellow=cell  cyan=equation_text_line"
    )


if __name__ == "__main__":
    main()