#!/usr/bin/env python3
"""
template.py — assemble the editable first-pass structural template from the spike's three signal
sources (extract structured.json + bands.json + furniture.json) into ONE round-trippable JSON the
human reviewer verifies AND edits before stage-2 generates the final template.

UI principle (user, 2026-06-07): directional LIMITS are draggable LINES (1-DOF, easier to drag);
object FOOTPRINTS are BOXES. So:
  * margins      -> four axis-locked LINES: left/right (x), top/bottom (y)
  * question/part bands -> horizontal LINES: start/end y
  * furniture / figures / tables -> BOXES (an object's footprint)

Every editable element carries {source: "auto"|"human", confirmed: bool} — the AI-suggestion seam.
Stage-2 must consume only confirmed elements (or a template marked confirmed at the top level).
Coordinates are PDF points, BOTTOM-LEFT origin (units in meta); the app maps to its own canvas.

Usage:
  python template.py --structured S.json --bands B.json --furniture F.json --pdf P.pdf --out T.json
"""
import json, argparse, datetime


def _line(edge, axis, value, scope, page=None):
    o = {"edge": edge, "axis": axis, "value": round(value, 1), "scope": scope,
         "source": "auto", "confirmed": False}
    if page is not None:
        o["page"] = page
    return o


def _furn_kind(it):
    """Best-guess label for a furniture box (human can rename). Position-based, BOTTOM-LEFT origin."""
    bb = it["bbox"]; cx = (bb["l"] + bb["r"]) / 2; cy = (bb["t"] + bb["b"]) / 2
    if it["kind"] == "picture":
        if cx > 430 and cy > 700:
            return "qr"
        if cy < 110:
            return "barcode"
        return "chrome_picture"
    if cy < 90:
        return "footer"
    if cy > 760:
        return "header_or_page_number"
    return "chrome_text"


def synthesize_part_box(part_band, content_x_band):
    """Return the one authoritative S5 part-box projection.

    Parts remain boxes in S5, but the box is a projection rather than intrinsic
    geometry: document content margins provide the x-extent and the part band
    provides y. The band end is already bounded by the next part in bands.py;
    the original label box remains a separate anchor for rendering/review.

    Coordinates stay in the first-pass PDF-point BOTTOMLEFT bbox shape.
    """
    if not content_x_band:
        return None
    try:
        x_left = content_x_band["x_left"]
        x_right = content_x_band["x_right"]
        y_start = part_band["y_start"]
        y_end = part_band["y_end"]
    except KeyError:
        return None
    return {
        "l": round(x_left, 1),
        "t": round(y_start, 1),
        "r": round(x_right, 1),
        "b": round(y_end, 1),
        "coord_origin": "BOTTOMLEFT",
    }


def build(structured, bands, furniture, pdf=None, page_roles=None):
    page_roles = page_roles or {}
    part_bbox = {p["label"]: p.get("bbox")
                 for q in structured.get("questions", []) for p in q["parts"]}
    cm = furniture.get("content_margins") or {}
    xband = cm.get("content_x_band") or {}
    per_pg_m = cm.get("per_page") or {}

    def margins_on(pg):
        r = page_roles.get(str(pg)) or page_roles.get(pg)
        return r.get("margins_enabled", True) if r else True

    # margins as axis-locked LINES — document-level left/right, per-page top/bottom. Per-page
    # top/bottom are omitted for pages with no content column (cover/blank) — the user's override.
    margins = []
    if "x_left" in xband:
        margins.append(_line("left", "x", xband["x_left"], "document"))
        margins.append(_line("right", "x", xband["x_right"], "document"))
    for pg, m in sorted(per_pg_m.items(), key=lambda kv: int(kv[0])):
        if not margins_on(int(pg)):
            continue
        margins.append(_line("top", "y", m["top"], "page", int(pg)))
        margins.append(_line("bottom", "y", m["bottom"], "page", int(pg)))

    # furniture + figures as BOXES, grouped by page
    furn_pg, fig_pg = {}, {}
    for it in furniture.get("items", []):
        pg = it["page"]
        if it.get("furniture"):
            furn_pg.setdefault(pg, []).append(
                {"box": it["bbox"], "kind": _furn_kind(it), "docling_label": it["label"],
                 "source": "auto", "confirmed": False})
        elif it["kind"] == "picture":
            fig_pg.setdefault(pg, []).append(
                {"box": it["bbox"], "source": "auto", "confirmed": False})

    tbl_pg = {}
    for t in structured.get("tables", []):
        if t.get("page"):
            tbl_pg.setdefault(t["page"], []).append(
                {"box": t.get("bbox"), "n_rows": t.get("n_rows"), "n_cols": t.get("n_cols"),
                 "table_source": t.get("source"), "source": "auto", "confirmed": False})

    # --- reconcile against recovered part labels -------------------------------------------
    # A part-label position is never furniture or a figure (the label wins), and a "figure" that
    # covers most of the content area is a Docling page-collapse artifact (the GPU sometimes flags
    # the whole page as one picture), not a real figure -> drop both. Fixes the Q1.7/Q1.9 clashes
    # and the full-page "figure" that was masking part labels.
    part_boxes_pg = {}
    for q in structured.get("questions", []):
        for p in q["parts"]:
            if p.get("bbox") and p.get("page"):
                part_boxes_pg.setdefault(p["page"], []).append(p["bbox"])

    def _inter(a, b):
        return not (a["r"] < b["l"] or b["r"] < a["l"] or a["t"] < b["b"] or b["t"] < a["b"])

    def _area(b):
        return max(0, b["r"] - b["l"]) * max(0, b["t"] - b["b"])

    for pg, items in list(furn_pg.items()):
        pls = part_boxes_pg.get(pg, [])
        furn_pg[pg] = [f for f in items if not (f.get("box") and any(_inter(f["box"], pl) for pl in pls))]
    for pg, items in list(fig_pg.items()):
        pls = part_boxes_pg.get(pg, [])
        m = per_pg_m.get(str(pg)) or per_pg_m.get(pg) or {}
        carea = ((m.get("right", 0) - m.get("left", 0)) * (m.get("top", 0) - m.get("bottom", 0))) or (595 * 842)
        fig_pg[pg] = [f for f in items if f.get("box")
                      and _area(f["box"]) <= 0.55 * carea               # not a full-page collapse
                      and not any(_inter(f["box"], pl) for pl in pls)]  # not clashing a part label

    pages = {}
    all_pg = (set(bands["pages"]) | {str(p) for p in furn_pg} | {str(p) for p in fig_pg}
              | {str(p) for p in page_roles})
    for pgs in sorted(all_pg, key=int):
        pg = int(pgs)
        pb = bands["pages"].get(pgs) or bands["pages"].get(pg) or {"main": [], "part": []}
        main = [{"question": m["question"], "y_start": m["y_start"], "y_end": m["y_end"],
                 "is_start": m.get("is_start", True),
                 "source": "auto", "confirmed": False} for m in pb["main"]]
        part = []
        for p in pb["part"]:
            part.append({
                "label": p["label"], "question": p["question"],
                "y_start": p["y_start"], "y_end": p["y_end"],
                "label_box": part_bbox.get(p["label"]),  # anchor, not the part extent
                "box": synthesize_part_box(p, xband),
                "marks": p.get("marks"),                 # parsed per-part marks (born-digital)
                "source": "auto", "confirmed": False,
            })
        pr = page_roles.get(pgs) or page_roles.get(pg) or {}
        pages[pgs] = {
            "role": pr.get("role", "question"),
            "role_source": pr.get("source", "default"), "role_confirmed": pr.get("confirmed", False),
            "margins_enabled": pr.get("margins_enabled", True),   # human-overridable
            "main_bands": main, "part_bands": part,
            "furniture": furn_pg.get(pg, []), "figures": fig_pg.get(pg, []),
            "tables": tbl_pg.get(pg, []),
        }

    return {
        "meta": {
            "schema": "exam-template/first-pass/v1",
            "board": structured.get("board"), "paper_code": structured.get("paper_code"),
            "source_pdf": pdf, "n_pages": furniture.get("n_pages"),
            "coord_origin": "BOTTOMLEFT", "units": "pdf_points",
            "generated_at": datetime.datetime.now().isoformat(timespec="seconds"),
            "ui_principle": "directional limits = draggable axis-locked lines; "
                            "object footprints = boxes",
            "confirmed": False, "confirmed_by": None, "confirmed_at": None,
        },
        "margins": margins,
        "pages": pages,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--structured", required=True)
    ap.add_argument("--bands", required=True)
    ap.add_argument("--furniture", required=True)
    ap.add_argument("--page-roles", dest="page_roles", help="page_roles.py JSON (roles + margin override)")
    ap.add_argument("--pdf")
    ap.add_argument("--out", default="results/template.json")
    a = ap.parse_args()
    roles = json.load(open(a.page_roles))["pages"] if a.page_roles else {}
    t = build(json.load(open(a.structured)), json.load(open(a.bands)),
              json.load(open(a.furniture)), a.pdf, roles)
    json.dump(t, open(a.out, "w"), indent=2)
    np = len(t["pages"])
    nm = sum(len(p["main_bands"]) for p in t["pages"].values())
    npt = sum(len(p["part_bands"]) for p in t["pages"].values())
    nf = sum(len(p["furniture"]) for p in t["pages"].values())
    ng = sum(len(p["figures"]) for p in t["pages"].values())
    print(f"template {t['meta']['paper_code']} ({t['meta']['board']}): {np} pages, "
          f"{len(t['margins'])} margin-lines, {nm} main-bands, {npt} part-bands, "
          f"{nf} furniture-boxes, {ng} figure-boxes")
    print(f"-> wrote {a.out}")


if __name__ == "__main__":
    main()
