#!/usr/bin/env python3
"""
overlay.py — human-viewable debug visualisation: draw the extractor's geometry over the rendered
exam page. Shows WHERE each question/part label was located and where Docling regions
(figures/tables/MCQ checkboxes) sit, so a reviewer can eyeball whether the structure landed in the
right place. This is the same geometry the exam-marker app uses to place regions on its canvas.

Coordinates: Docling/RapidOCR bboxes are PDF points with a BOTTOM-LEFT origin. We render the page
at DPI D (scale = D/72) and flip y against the rendered image height, so we never need the page's
point-height explicitly: y_top_px = H_px - t*scale.

With --docling, also draws every raw Docling text block (the body/question content the thin
extractor model discards) so a reviewer can see the FULL detection, not just what we persist.
Granite tables carry cells but no coordinates; we derive their box by locating the cell-texts in
the Docling text layer (content+geometry fusion).

Usage:
  python scripts/overlay.py <structured.json> <source_pdf> [--pages 3,4,5] [--dpi 150] [--out DIR]
  python scripts/overlay.py <structured.json> <pdf> --docling results/E_tess_full.json --pages 5
"""
import os, sys, json, re, argparse, subprocess, tempfile
from PIL import Image, ImageDraw, ImageFont

PART_COLOR = (211, 47, 47)                          # red — question/part labels
BODY_COLOR = (150, 150, 150)                         # grey — raw Docling body-text blocks (--docling)
GRANITE_COLOR = (0, 150, 136)                        # teal — Granite table (geometry derived from cells)
REGION_COLORS = {                                   # docling region taxonomy -> colour
    "context_figure": (25, 118, 210),               # blue
    "context_data": (56, 142, 60),                  # green (tables)
    "context_caption": (123, 31, 162),              # purple
    "mcq_option": (245, 124, 0),                     # orange (checkboxes)
}


def _norm(s):
    return re.sub(r"[^a-z0-9]", "", (s or "").lower())


def docling_texts_by_page(doc):
    """All raw Docling text items -> {page: [(bbox, text, label)]}. The body content we discard."""
    out = {}
    for t in doc.get("texts", []):
        prov = t.get("prov") or []
        bb = prov[0].get("bbox") if prov else None
        pg = prov[0].get("page_no") if prov else None
        if bb and pg:
            out.setdefault(pg, []).append((bb, t.get("text") or "", t.get("label") or "text"))
    return out


def derive_table_bbox(grid, page_texts):
    """Granite tables have cells but no coordinates. Locate the cell-texts in the Docling text
    layer and union their bboxes -> the table's on-page extent.

    Two traps (seen on physics p5): (1) border/maths glyphs ('|','+') normalise to '' and an
    empty string is a substring of everything; (2) cell WORDS recur in nearby content — the rock
    names reappear in the MCQ options below the table ('Basalt or chalk'), far left and lower.
    So we match only blocks whose normalised text is CONTAINED IN a cell (keeps fragments like
    '2.90'/'Type', rejects the longer 'basaltorchalk'), require length >= 2, then keep the
    dominant vertical cluster to drop any stray cell-word elsewhere on the page."""
    import statistics
    cells = {c for c in (_norm(x) for row in grid for x in row) if len(c) > 1}
    hit = [bb for bb, txt, _ in page_texts
           if len(_norm(txt)) > 1 and any(_norm(txt) in c for c in cells)]
    if len(hit) < 3:
        return None
    med = statistics.median(sorted((b["t"] + b["b"]) / 2 for b in hit))
    hit = [b for b in hit if abs((b["t"] + b["b"]) / 2 - med) <= 120]   # table band only
    return {"l": min(b["l"] for b in hit), "r": max(b["r"] for b in hit),
            "t": max(b["t"] for b in hit), "b": min(b["b"] for b in hit)}


def _font(sz):
    for p in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
              "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
        if os.path.exists(p):
            return ImageFont.truetype(p, sz)
    return ImageFont.load_default()


MAIN_LINE = (25, 118, 210)                          # blue — main-question y-markers
PART_LINE = (211, 47, 47)                            # red — part y-markers


def _hline(draw, y_pdf, scale, H, W, color, label, width, font, dashed=False, inset=0):
    """Full-width horizontal marker line at a PDF-point y (BOTTOM-LEFT origin)."""
    y = H - y_pdf * scale
    if dashed:
        x = inset
        while x < W:
            draw.line([x, y, min(x + 9, W), y], fill=color, width=width); x += 16
    else:
        draw.line([inset, y, W, y], fill=color, width=width)
    if label:
        tw = draw.textlength(label, font=font)
        draw.rectangle([inset, y - 16, inset + tw + 6, y], fill=color)
        draw.text((inset + 3, y - 15), label, fill=(255, 255, 255), font=font)


def _rect(draw, bb, scale, H, color, label, width=3, font=None):
    """Draw one bbox (BOTTOM-LEFT origin -> image space) + its label."""
    x0, x1 = bb["l"] * scale, bb["r"] * scale
    y0, y1 = H - bb["t"] * scale, H - bb["b"] * scale      # t is the higher edge -> smaller y_px
    draw.rectangle([x0, y0, x1, y1], outline=color, width=width)
    if label:
        tw = draw.textlength(label, font=font)
        draw.rectangle([x0, y0 - 17, x0 + tw + 6, y0], fill=color)
        draw.text((x0 + 3, y0 - 16), label, fill=(255, 255, 255), font=font)


def draw_template(draw, tpl, pg, scale, H, W, font):
    """Render the editable template for one page: margins/bands as LINES, footprints as BOXES.
    A confirmed element is drawn solid; an unconfirmed (auto) suggestion is drawn dashed."""
    MARGIN, MAIN, PART = (0, 150, 136), (25, 118, 210), (211, 47, 47)
    page = tpl["pages"].get(str(pg)) or tpl["pages"].get(pg) or {}
    # role banner (top-left); margins suppressed entirely on no-margin pages (cover/blank)
    role = page.get("role", "question")
    draw.rectangle([0, 0, 8 + len(role) * 8, 16], fill=(70, 70, 70))
    draw.text((4, 1), f"role:{role}", fill=(255, 255, 255), font=font)
    margins_on = page.get("margins_enabled", True)
    # margins: axis-locked lines (document scope on every page + this page's page-scope lines)
    for m in (tpl.get("margins", []) if margins_on else []):
        if m["scope"] == "page" and m.get("page") != pg:
            continue
        solid = m.get("confirmed")
        if m["axis"] == "x":
            x = m["value"] * scale
            draw.line([x, 0, x, H], fill=MARGIN, width=2) if solid else _dash_v(draw, x, 0, H, MARGIN, 2)
        else:
            y = H - m["value"] * scale
            draw.line([0, y, W, y], fill=MARGIN, width=2) if solid else _dash_h(draw, 0, W, y, MARGIN, 2)
    for m in page.get("main_bands", []):
        if not m.get("is_start", True):          # continuation page: no spurious second "start" line
            continue
        _hline(draw, m["y_start"], scale, H, W, MAIN, f"Q{m['question']}", 3, font,
               dashed=not m.get("confirmed"))
    for p in page.get("part_bands", []):
        _hline(draw, p["y_start"], scale, H, W, PART, p["label"], 2, font, inset=90,
               dashed=not p.get("confirmed"))
    for f in page.get("furniture", []):
        if f.get("box"):
            _rect(draw, f["box"], scale, H, (130, 130, 130), f"furniture:{f.get('kind','')}", 2, font)
    for g in page.get("figures", []):
        if g.get("box"):
            _rect(draw, g["box"], scale, H, (56, 142, 60), "figure", 3, font)
    for t in page.get("tables", []):
        if t.get("box"):
            _rect(draw, t["box"], scale, H, (0, 150, 136),
                  f"table {t.get('n_rows')}x{t.get('n_cols')}", 3, font)


def render_page(pdf, pg, dpi, td):
    """Render page `pg` and return an image in DOCLING's coordinate space. Docling reports bbox
    relative to the CropBox, but pdftoppm renders the MediaBox — when CropBox != MediaBox (e.g. the
    Edexcel 1MA1 papers: media 652x899, crop inset 28.35pt) that mismatch magnifies + shifts every
    overlaid shape toward a corner. Fix: crop the rendered image to the CropBox so it matches Docling.
    No-op when CropBox == MediaBox (h556) or when poppler already rendered the CropBox."""
    base = os.path.join(td, f"p{pg}")
    subprocess.run(["pdftoppm", "-png", "-r", str(dpi), "-f", str(pg), "-l", str(pg), pdf, base],
                   check=True)
    png = next(p for p in (f"{base}-{pg:02d}.png", f"{base}-{pg}.png", f"{base}-{pg:03d}.png")
               if os.path.exists(p))
    img = Image.open(png).convert("RGB")
    try:
        import pypdf
        page = pypdf.PdfReader(pdf).pages[pg - 1]
        mb, cb = page.mediabox, page.cropbox
        scale = dpi / 72.0
        mbl, mbt = float(mb.left), float(mb.top)
        dcrop = any(abs(a - b) > 0.5 for a, b in
                    ((cb.left, mb.left), (cb.bottom, mb.bottom), (cb.right, mb.right), (cb.top, mb.top)))
        rendered_mediabox = abs(img.width - (float(mb.right) - mbl) * scale) < 3
        if dcrop and rendered_mediabox:
            img = img.crop((round((float(cb.left) - mbl) * scale), round((mbt - float(cb.top)) * scale),
                            round((float(cb.right) - mbl) * scale), round((mbt - float(cb.bottom)) * scale)))
    except Exception:
        pass
    return img


def _dash_v(draw, x, y0, y1, color, w):
    y = y0
    while y < y1:
        draw.line([x, y, x, min(y + 9, y1)], fill=color, width=w); y += 16


def _dash_h(draw, x0, x1, y, color, w):
    x = x0
    while x < x1:
        draw.line([x, y, min(x + 9, x1), y], fill=color, width=w); x += 16


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("structured"); ap.add_argument("pdf")
    ap.add_argument("--docling", help="raw Docling doc JSON: also draw every body-text block "
                                      "(the content the thin model discards) + derive Granite-table boxes")
    ap.add_argument("--bands", help="bands.py JSON: draw main-question + part start/end y-marker lines")
    ap.add_argument("--furniture", help="furniture.py JSON: mark recurring furniture vs real figures "
                                        "+ draw the content x-margins")
    ap.add_argument("--template", help="template.py JSON: render the editable first-pass template "
                                       "(margins+bands as lines, furniture/figures as boxes). "
                                       "When set, draws ONLY the template (the human-review view).")
    ap.add_argument("--pages", help="comma list, e.g. 3,4,5 (default: all pages with geometry)")
    ap.add_argument("--dpi", type=int, default=150)
    ap.add_argument("--out", default="results/overlay")
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)
    scale = a.dpi / 72.0
    font = _font(14)

    res = json.load(open(a.structured))
    doc_texts = docling_texts_by_page(json.load(open(a.docling))) if a.docling else {}
    bands = json.load(open(a.bands))["pages"] if a.bands else {}
    furn = json.load(open(a.furniture)) if a.furniture else None
    tpl = json.load(open(a.template)) if a.template else None
    # gather geometry by page
    parts_by_pg, regions_by_pg = {}, {}
    for q in res.get("questions", []):
        for p in q["parts"]:
            if p.get("bbox") and p.get("page"):
                parts_by_pg.setdefault(p["page"], []).append((p["label"], p["bbox"]))
    for r in res.get("regions", []):
        if r.get("bbox") and r.get("page"):
            regions_by_pg.setdefault(r["page"], []).append((r["type"], r["bbox"]))
    # tables: standard ones carry a bbox; Granite ones don't -> derive from the text layer
    tables_by_pg = {}
    for t in res.get("tables", []):
        pg = t.get("page")
        if not pg:
            continue
        bb = t.get("bbox") or (derive_table_bbox(t.get("grid", []), doc_texts.get(pg, []))
                               if a.docling else None)
        if bb:
            tables_by_pg.setdefault(pg, []).append(
                (f"table {t.get('source','')} {t.get('n_rows')}x{t.get('n_cols')}", bb))

    want = ([int(x) for x in a.pages.split(",")] if a.pages
            else (sorted(int(p) for p in tpl["pages"]) if tpl
                  else sorted(set(parts_by_pg) | set(regions_by_pg) | set(doc_texts))))
    if not want:
        sys.exit("no bbox geometry in this result (born-digital text path carries no geometry; "
                 "use an OCR/rapid-path structured.json)")

    written = []
    with tempfile.TemporaryDirectory() as td:
        for pg in want:
            img = render_page(a.pdf, pg, a.dpi, td)
            H = img.height
            draw = ImageDraw.Draw(img)
            if tpl:                          # template-only render = the human-review view
                draw_template(draw, tpl, pg, scale, H, img.width, font)
                out = os.path.join(a.out, f"p{pg:02d}.png")
                img.save(out); written.append(out)
                pgd = tpl["pages"].get(str(pg), {})
                print(f"p{pg}: template — {len(pgd.get('main_bands',[]))} main, "
                      f"{len(pgd.get('part_bands',[]))} part, {len(pgd.get('furniture',[]))} furn, "
                      f"{len(pgd.get('figures',[]))} fig -> {out}")
                continue
            # layer 0: raw Docling body-text blocks (faint, no label) — the discarded content
            for bb, txt, lab in doc_texts.get(pg, []):
                _rect(draw, bb, scale, H, BODY_COLOR, None, 1, font)
            # layer 1: taxonomy regions
            for typ, bb in regions_by_pg.get(pg, []):
                _rect(draw, bb, scale, H, REGION_COLORS.get(typ, (120, 120, 120)), typ, 2, font)
            # layer 2: tables (Granite-derived boxes in teal)
            for lab, bb in tables_by_pg.get(pg, []):
                _rect(draw, bb, scale, H, GRANITE_COLOR, lab, 3, font)
            # layer 3: part labels on top
            for lab, bb in parts_by_pg.get(pg, []):
                _rect(draw, bb, scale, H, PART_COLOR, lab, 3, font)
            # layer 4: band y-marker lines (main-question = blue, part = red dashed; end = dashed)
            pb = bands.get(str(pg)) or bands.get(pg)
            nb = 0
            if pb:
                W = img.width
                for m in pb["main"]:
                    if not m.get("is_start", True):     # skip continuation-page duplicate
                        continue
                    _hline(draw, m["y_start"], scale, H, W, MAIN_LINE,
                           f"Q{m['question']} ▸ start", 3, font); nb += 1
                    _hline(draw, m["y_end"], scale, H, W, MAIN_LINE, None, 1, font, dashed=True)
                for p in pb["part"]:
                    _hline(draw, p["y_start"], scale, H, W, PART_LINE,
                           f"{p['label']} start", 2, font, inset=90); nb += 1
            # layer 5: furniture mask — green=real figure, grey=masked furniture; + content margins
            if furn:
                W = img.width
                for it in furn["items"]:
                    if it["page"] != pg or it["kind"] != "picture":
                        continue
                    if it["furniture"]:
                        _rect(draw, it["bbox"], scale, H, (130, 130, 130), "furniture", 2, font)
                    else:
                        _rect(draw, it["bbox"], scale, H, (56, 142, 60), "figure ✓", 3, font)
                band = (furn.get("content_margins") or {}).get("content_x_band")
                if band:
                    for xk in ("x_left", "x_right"):
                        x = band[xk] * scale
                        draw.line([x, 0, x, H], fill=(0, 150, 136), width=2)
            out = os.path.join(a.out, f"p{pg:02d}.png")
            img.save(out); written.append(out)
            print(f"p{pg}: {len(parts_by_pg.get(pg,[]))} part-labels, "
                  f"{len(regions_by_pg.get(pg,[]))} regions, {len(tables_by_pg.get(pg,[]))} tables, "
                  f"{len(doc_texts.get(pg,[]))} body-text blocks, {nb} band-lines -> {out}")
    print(f"-> {len(written)} page(s) in {a.out}/")


if __name__ == "__main__":
    main()
