#!/usr/bin/env python3
"""Cut one crop per mapped row, with the element framed in red, for the Screenshot column.

A crop keeps the whole frame width and a strip of --pad pixels over and under
the element, which gets a red outline. Output lands in
keyshots/keyshot__NN__<key>.png, NN being the row's index in mapping.rows
counted from 1. A row whose element cannot be found on screen is reported as
missing and gets no picture; crops are never guessed. Node matching and the
Figma calls are shared with render-overlay.py, loaded from this folder.
"""
import argparse
import importlib.util
import json
import os
import re
import subprocess
import sys

SCRIPT_DIR = os.path.split(os.path.abspath(__file__))[0]
BOX_COLOUR = "#E81932"
ID_FIELD = "cmsNodeId"
INSET = 6
UNSAFE = re.compile(r"[\\/:*?\"<>|\s]+")
TRAILING_MARKER = re.compile(r"\s*\([^()]*\)\s*$")


def load_overlay():
    path = os.path.join(SCRIPT_DIR, "render-overlay.py")
    try:
        spec = importlib.util.spec_from_file_location("render_overlay", path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
    except Exception as exc:
        sys.exit(f"cannot load {path}: {exc}")
    return module


ov = load_overlay()


def shot_name(row_number, key):
    clean = UNSAFE.sub("-", TRAILING_MARKER.sub("", key).strip()).strip("-")
    return f"keyshot__{row_number:02d}__{clean}"


def spec_pages(paths):
    pages = []
    for path in paths:
        spec, loaded = ov.load_spec_pages(path)
        if not loaded:
            sys.exit(f"--spec lists no pages: {path}")
        pages.extend(loaded)
    return pages


def visible_text_nodes(document):
    frame = document.get("absoluteBoundingBox") or {}
    fx, fy = frame.get("x", 0), frame.get("y", 0)
    found = []
    pending = [(document, True, 1.0)]
    while pending:
        node, shown, alpha = pending.pop()
        shown = shown and node.get("visible", True) is not False
        alpha = alpha * float(node.get("opacity", 1.0))
        if not shown or alpha <= 0:
            continue
        box = node.get("absoluteBoundingBox")
        text = node.get("characters")
        is_text = node.get("type") == "TEXT"
        if is_text and box and text and text.strip():
            found.append({
                "id": node.get("id"),
                "characters": text,
                "x": round(box["x"] - fx),
                "y": round(box["y"] - fy),
                "w": round(box["width"]),
                "h": round(box["height"]),
            })
        pending.extend((child, shown, alpha) for child in reversed(node.get("children") or []))
    return found


def rest_pages(file_key, node_ids, scale, token, ca_file):
    pages = []
    for raw in node_ids:
        node_id = ov.canon_id(raw)
        document = ov.fetch_node(file_key, node_id, token, ca_file)
        if not document:
            print(f"# warning: {file_key} has no frame {node_id}; skipping it", file=sys.stderr)
            continue
        image = ov.render_image(file_key, node_id, scale, token, ca_file)
        if not image:
            print(f"# warning: Figma returned no render of {node_id}; skipping it", file=sys.stderr)
            continue
        box = document.get("absoluteBoundingBox") or {}
        pages.append({
            "nodeId": node_id,
            "name": document.get("name"),
            "image": image,
            "frame": {"w": round(box.get("width", 0)), "h": round(box.get("height", 0))},
            "nodes": visible_text_nodes(document),
        })
    if not pages:
        sys.exit("none of the frames could be rendered")
    return pages


def assign(rows, pages):
    index = ov.build_index(rows)
    hits = {}
    for page in pages:
        ordered = sorted(page.get("nodes") or [], key=lambda item: float(item.get("y", 0)))
        for node in ordered:
            position = ov.match_row(node, index)
            if position is not None:
                hits.setdefault(position, (page, node))
    return hits


def crop_html(page, node, pad):
    fw, fh = (float(page["frame"][side]) for side in ("w", "h"))
    x, y, w, h = (float(node.get(axis, 0)) for axis in ("x", "y", "w", "h"))
    top = max(0.0, y - pad)
    height = max(1.0, min(fh, y + h + pad) - top)
    pin_left = "left:0;"
    box_class = "box"
    style = (
        f"html,body{{margin:0;padding:0;}}body{{width:{fw:.0f}px;height:{height:.0f}px;position:relative;overflow:hidden;}}"
        f"img{{position:absolute;{pin_left}top:{-top:.0f}px;width:{fw:.0f}px;height:{fh:.0f}px;}}"
        f".{box_class}{{position:absolute;left:{x - INSET:.0f}px;top:{y - top - INSET:.0f}px;width:{w + 2 * INSET:.0f}px;"
        f"height:{h + 2 * INSET:.0f}px;border:3px solid {BOX_COLOUR};border-radius:6px;box-sizing:border-box;}}"
    )
    document = (
        f"{ov.PAGE_OPEN}{ov.META_UTF8}<style>{style}</style></head>"
        f'<body><img src="{page["image"]}"><div class="{box_class}"></div></body></html>'
    )
    return document, fw, height


def main():
    cli = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    cli.add_argument("--mapping", required=True, help="the reuse-map JSON")
    cli.add_argument("--spec", action="append", help="geometry spec for one screen state; may be repeated")
    cli.add_argument("--file", help="Figma file key; falls back to figmaFileKey in the mapping")
    cli.add_argument("--nodes", help="frame node-ids separated by commas; falls back to figmaNodes")
    cli.add_argument("--pad", default=150, type=int, help="pixels of surrounding context kept over and under the element")
    cli.add_argument("--scale", default=2, type=int)
    ov.add_token_flag(cli)
    cli.add_argument("--ca-file", help="PEM roots to trust when a proxy intercepts TLS")
    cli.add_argument("--out", dest="target_dir", default=".", help="folder for the manifest; the crops go into <out>/keyshots/")
    cli.add_argument("--slug", help="stem of the manifest name; derived from the screen name when omitted")
    opts = cli.parse_args()
    ov.reject_token_flag(cli, opts, ov.FIGMA_TOKEN_HINT)

    with open(opts.mapping, encoding="utf-8") as stream:
        mapping = json.load(stream)
    rows = mapping.get("rows") or []
    stem = opts.slug or ov.slug(mapping.get("screen"))

    if opts.spec:
        pages = spec_pages(opts.spec)
    else:
        file_key = opts.file or mapping.get("figmaFileKey")
        nodes = ov.split_nodes(opts.nodes, mapping)
        if not (file_key and nodes):
            cli.error("need a Figma file key and frame node-ids (flags or mapping), or --spec")
        token = ov.figma_token()
        if not token:
            sys.exit("No Figma token found (FIGMA_ACCESS_TOKEN, FIGMA_TOKEN or keychain FIGMA_ACCESS_TOKEN). "
                     "Without one, build --spec files from the screenshots instead.")
        pages = rest_pages(file_key, nodes, opts.scale, token, opts.ca_file)

    crops_dir = os.path.join(opts.target_dir, "keyshots")
    os.makedirs(crops_dir, exist_ok=True)
    chrome = ov.find_chrome()
    hits = assign(rows, pages)
    shots, missing = [], []
    failed = 0

    for number, row in enumerate(rows, 1):
        key = str(row.get("newKey") or "").strip()
        if not key:
            continue
        if number - 1 not in hits:
            missing.append({"row": number, "key": key, "reason": f"no visible node matched ({ID_FIELD}/value)"})
            continue
        page, node = hits[number - 1]
        document, width, height = crop_html(page, node, opts.pad)
        base = shot_name(number, key)
        entry = {
            "row": number,
            "key": TRAILING_MARKER.sub("", key),
            "element": row.get("element"),
            "pageNode": page.get("nodeId"),
            "cropW": int(round(width)),
            "cropH": int(round(height)),
        }
        if not chrome:
            with open(os.path.join(crops_dir, base + ".html"), "w", encoding="utf-8") as sink:
                sink.write(document)
            entry.update(file=f"keyshots/{base}.html", fallback="html")
            shots.append(entry)
            continue
        try:
            ov.chrome_screenshot(chrome, document, os.path.join(crops_dir, base + ".png"), width, height, opts.scale)
        except (RuntimeError, subprocess.SubprocessError) as exc:
            failed += 1
            print(f"# row {number} {key}: {exc}", file=sys.stderr)
            continue
        entry["file"] = f"keyshots/{base}.png"
        shots.append(entry)

    manifest_name = f"{stem}.keyshots.manifest.json"
    manifest = {
        "screen": mapping.get("screen"),
        "slug": stem,
        "pad": opts.pad,
        "scale": opts.scale,
        "dir": "keyshots",
        "shots": shots,
        "missing": missing,
    }
    with open(os.path.join(opts.target_dir, manifest_name), "w", encoding="utf-8") as sink:
        json.dump(manifest, sink, indent=2, ensure_ascii=False)
        sink.write("\n")

    if not chrome:
        print("# Chrome, Chromium and Edge are all missing; the crops were saved as .html", file=sys.stderr)
    print(f"# keyshots: {len(shots)} written, {len(missing)} unmatched, {failed} failed", file=sys.stderr)
    print(f'# point the mapping\'s top-level "keyshots" field at "{manifest_name}"', file=sys.stderr)


if __name__ == "__main__":
    main()