#!/usr/bin/env python3
"""Draw the overlay that ties each localization key to its place on one screen.

The screen image sits in the middle; every keyed element gets a numbered card
at the side and a dashed line leading to it. A card's number equals the index,
counted from 1, of its matching entry in mapping.rows. Boxes are read from
Figma over REST, or taken from a --spec file, and headless Chrome turns the
HTML page into a PNG. The same input always gives the same picture; nothing is
generated by a model.
"""
import argparse
import base64
import html
import json
import os
import re
import shutil
import ssl
import subprocess
import sys
import tempfile
import time
import unicodedata
import urllib.error
import urllib.parse
import urllib.request

from _shared import add_token_flag, reject_token_flag, slug

FIGMA_TOKEN_HINT = "export FIGMA_ACCESS_TOKEN (or FIGMA_TOKEN), or keep it in the macOS keychain under FIGMA_ACCESS_TOKEN"

API = "https://api.figma.com/v1"
SEP = "[:\\-\u2013]"
TR_RE = re.compile("^(?:tr|tr[kc]e|t[\u00fcu]rk[\u00e7c]e)\\s*" + SEP + "\\s*(.+)$", re.IGNORECASE)
EN_RE = re.compile("^(?:en|eng|ing|[i\u0130]ngilizce)\\s*" + SEP + "\\s*(.+)$", re.IGNORECASE)

CARD_W = 320
COLUMN_X = 24
CHANNEL = 190
TOP = 60
CARD_H = 116
GAP = 14
RIGHT_MARGIN = 40
PINK = "#EC1E91"
GRAY = "#9aa0a6"
BACKGROUND = "#1e1e1e"
LABELS = {
    "en": {"cms": "CMS copy", "pending": "Awaiting copy", "unmatched": "(unmatched)"},
    "tr": {"cms": "\u0130\u00e7erik (CMS)", "pending": "\u00d6neri bekleniyor", "unmatched": "(e\u015fle\u015fmedi)"},
}

PAGE_OPEN = "<!doctype html><html><head>"
META_UTF8 = '<meta charset="utf-8">'


def mac_app(bundle):
    return f"/Applications/{bundle}.app/Contents/MacOS/{bundle}"


MAC_CHROMES = [mac_app(name) for name in ("Google Chrome", "Chromium", "Microsoft Edge")]
PATH_CHROMES = (
    [f"google-chrome{tail}" for tail in ("", "-stable")]
    + [f"chromium{tail}" for tail in ("", "-browser")]
    + ["chrome"]
)


URL_NODE_ID = re.compile(r"^\d+-\d+$")


def canon_id(node_id):
    text = str(node_id or "").strip()
    return text.replace("-", ":") if URL_NODE_ID.match(text) else text


def split_lang(text):
    picked = {"tr": None, "en": None}
    loose = []
    for line in filter(None, (piece.strip() for piece in str(text or "").splitlines())):
        tr_hit, en_hit = TR_RE.match(line), EN_RE.match(line)
        if tr_hit and picked["tr"] is None:
            picked["tr"] = tr_hit.group(1).strip()
        elif en_hit and picked["en"] is None:
            picked["en"] = en_hit.group(1).strip()
        else:
            loose.append(line)
    if picked["tr"] is None and picked["en"] is None:
        padded = loose + [None, None]
        style = {0: "empty", 1: "single"}.get(len(loose), "two-line")
        return {"tr": padded[0], "en": padded[1], "mode": style}
    spare = iter(loose)
    for lang in ("tr", "en"):
        if picked[lang] is None:
            picked[lang] = next(spare, None)
    return dict(picked, mode="prefixed")


def norm(value):
    text = unicodedata.normalize("NFC", str(value or "")).strip()
    text = re.sub(r"\s+", " ", text)
    text = re.sub(r"[.:\u2026]+$", "", text).strip()
    return text.lower()


def build_index(rows):
    by_id, by_value = {}, {}
    for position, entry in enumerate(rows):
        if entry.get("cmsNodeId"):
            by_id.setdefault(canon_id(entry["cmsNodeId"]), position)
        texts = [(entry.get(block) or {}).get(lang) for block in ("new", "cms", "legacy") for lang in ("tr", "en")]
        texts.append(entry.get("element"))
        for text in texts:
            normalised = norm(text)
            if normalised:
                by_value.setdefault(normalised, position)
    return {"id": by_id, "value": by_value}


def match_row(node, index):
    node_id = canon_id(node.get("id"))
    if node_id in index["id"]:
        return index["id"][node_id]
    for field in ("tr", "characters"):
        key = norm(node.get(field))
        if key and key in index["value"]:
            return index["value"][key]
    return None


def keychain_secret(service):
    try:
        lookup = subprocess.run(
            ["security", "find-generic-password", "-s", service, "-w"],
            capture_output=True, text=True, timeout=10,
        )
    except (OSError, subprocess.SubprocessError):
        return None
    secret = lookup.stdout.strip()
    return secret if lookup.returncode == 0 and secret else None


def figma_token():
    from_env = next((os.environ[v] for v in ("FIGMA_ACCESS_TOKEN", "FIGMA_TOKEN") if os.environ.get(v)), None)
    return from_env or keychain_secret("FIGMA_ACCESS_TOKEN")


def _open(url, headers, ca_file, timeout):
    tls = ssl.create_default_context(cafile=ca_file) if ca_file else None
    try:
        with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=timeout, context=tls) as reply:
            return reply.read()
    except urllib.error.HTTPError as exc:
        if exc.code == 429:
            sys.exit("Figma is rate limiting (HTTP 429). Measure the boxes on the screenshot and supply them through --spec.")
        snippet = exc.read()[:200].decode("utf-8", "replace")
        sys.exit(f"Figma HTTP {exc.code} for {url}: {snippet}")
    except urllib.error.URLError as exc:
        sys.exit(f"Could not reach Figma: {exc.reason}")


def http_json(url, token, ca_file, timeout=60):
    return json.loads(_open(url, {"X-Figma-Token": token}, ca_file, timeout).decode("utf-8"))


def http_bytes(url, ca_file, timeout=120):
    return _open(url, {}, ca_file, timeout)


def as_data_uri(png_bytes):
    return "data:image/png;base64," + base64.b64encode(png_bytes).decode("ascii")


def fetch_node(file_key, node_id, token, ca_file):
    ids = urllib.parse.quote(node_id, safe="")
    reply = http_json(f"{API}/files/{file_key}/nodes?ids={ids}", token, ca_file)
    entry = (reply.get("nodes") or {}).get(node_id) or {}
    return entry.get("document")


def render_image(file_key, node_id, scale, token, ca_file):
    ids = urllib.parse.quote(node_id, safe="")
    reply = http_json(f"{API}/images/{file_key}?ids={ids}&format=png&scale={scale}", token, ca_file)
    link = (reply.get("images") or {}).get(node_id)
    return as_data_uri(http_bytes(link, ca_file)) if link else None


def walk(root):
    stack = [root]
    while stack:
        node = stack.pop()
        if isinstance(node, dict):
            yield node
            stack.extend(reversed(node.get("children") or []))


def annotation_copy(node):
    for note in node.get("annotations") or []:
        text = note.get("label") or note.get("labelMarkdown") if isinstance(note, dict) else note
        parsed = split_lang(text)
        if parsed["tr"] or parsed["en"]:
            return parsed
    return None


def image_source(value, base_dir):
    if not value or str(value).startswith("data:"):
        return value
    path = os.path.join(base_dir, value) if not os.path.isabs(value) else value
    if os.path.isfile(path):
        with open(path, "rb") as stream:
            return as_data_uri(stream.read())
    sys.exit(f"the --spec image does not exist: {path}")


def page_from_rest(file_key, node_id, scale, token, ca_file):
    document = fetch_node(file_key, node_id, token, ca_file)
    if not document:
        sys.exit(f"Figma file {file_key} has no node {node_id}")
    box = document.get("absoluteBoundingBox") or {}
    fx, fy = box.get("x", 0), box.get("y", 0)
    nodes = []
    for node in walk(document):
        bbox = node.get("absoluteBoundingBox")
        if not bbox or node is document:
            continue
        copy = annotation_copy(node)
        text = node.get("characters") if node.get("type") == "TEXT" else None
        if not (text and text.strip()) and not copy:
            continue
        nodes.append({
            "id": node.get("id"),
            "characters": text,
            "tr": copy["tr"] if copy else None,
            "en": copy["en"] if copy else None,
            "x": round(bbox["x"] - fx),
            "y": round(bbox["y"] - fy),
            "w": round(bbox["width"]),
            "h": round(bbox["height"]),
            "side": "left",
        })
    image = render_image(file_key, node_id, scale, token, ca_file)
    if not image:
        sys.exit(f"Figma returned no image for {node_id}")
    return {
        "nodeId": node_id,
        "name": document.get("name"),
        "image": image,
        "frame": {"w": round(box.get("width", 0)), "h": round(box.get("height", 0))},
        "nodes": nodes,
    }


def build_cards(page, rows, labels):
    index = build_index(rows)
    cards = []
    taken = set()
    extra_number = len(rows)
    for node in page.get("nodes") or []:
        position = match_row(node, index)
        if position is not None and position in taken:
            continue
        entry = rows[position] if position is not None else None
        cms = (entry.get("cms") if entry else {}) or {}
        tr = node.get("tr") or cms.get("tr")
        en = node.get("en") or cms.get("en")
        new_key = entry.get("newKey") if entry else None
        if not (new_key or tr or en):
            continue
        if position is None:
            extra_number += 1
            number = extra_number
        else:
            taken.add(position)
            number = position + 1
        box = {axis: float(node.get(axis, 0)) for axis in ("x", "y", "w", "h")}
        cards.append({
            "num": number,
            "key": new_key or labels["unmatched"],
            "tr": tr,
            "en": en,
            "text": node.get("characters") or "",
            "annotated": bool(tr or en),
            "side": "right" if node.get("side") == "right" else "left",
            **box,
        })
    return cards


def centre_y(card):
    return card["y"] + card["h"] / 2


def stack_side(cards):
    cards.sort(key=centre_y)
    bottom = TOP - GAP
    for card in cards:
        target = TOP + card["y"] + card["h"] / 2 - CARD_H / 2
        card["cardY"] = max(target, bottom + GAP, TOP)
        bottom = card["cardY"] + CARD_H
    return bottom


def chip(kind, colour, content):
    return f'<span class="{kind}" style="background:{colour}">{content}</span>'


def card_html(card, left, labels):
    colour = PINK if card["annotated"] else GRAY
    pill = labels["cms"] if card["annotated"] else labels["pending"]
    if card["annotated"]:
        body = "".join(
            f'<div class="line"><b>{label}</b> {html.escape(card[lang] or "-")}</div>'
            for label, lang in (("TR", "tr"), ("EN", "en"))
        )
    else:
        body = f'<div class="line muted">{html.escape(card["text"] or "-")}</div>'
    head = chip("num", colour, card["num"]) + chip("pill", colour, html.escape(pill))
    return (
        f'<div class="card" style="left:{left}px;top:{card["cardY"]:.0f}px;border-color:{colour}">'
        f'<div class="head">{head}</div>'
        f'<div class="key">{html.escape(card["key"])}</div>{body}</div>'
    )


def connector(card, lane, lane_count, start_x, target_x, from_left):
    colour = PINK if card["annotated"] else GRAY
    step = (CHANNEL - 40) / max(lane_count, 1)
    bend = start_x + 20 + step * lane if from_left else start_x - 20 - step * lane
    card_mid = card["cardY"] + CARD_H / 2
    dot_y = TOP + centre_y(card)
    x, y = f"{target_x:.0f}", f"{dot_y:.0f}"
    label_y = f"{dot_y + 4:.0f}"
    route = f"{start_x:.0f},{card_mid:.0f} {bend:.0f},{card_mid:.0f} {bend:.0f},{y} {x},{y}"
    return (
        f'<polyline points="{route}" fill="none" stroke="{colour}" stroke-width="2" stroke-dasharray="6 4"/>'
        f'<circle cx="{x}" cy="{y}" r="11" fill="{colour}"/>'
        f'<text x="{x}" y="{label_y}" text-anchor="middle" font-size="12" font-weight="700" fill="#fff">{card["num"]}</text>'
    )


def build_html(page, cards, labels):
    fw, fh = page["frame"]["w"], page["frame"]["h"]
    sides = {"left": [], "right": []}
    for card in cards:
        sides[card["side"]].append(card)
    screen_x = COLUMN_X + CARD_W + CHANNEL
    right_x = screen_x + fw + CHANNEL
    width = right_x + CARD_W + RIGHT_MARGIN if sides["right"] else screen_x + fw + RIGHT_MARGIN
    height = int(max(TOP + fh, stack_side(sides["left"]), stack_side(sides["right"])) + TOP)

    boxes, wires = [], []
    for card_x, group in ((COLUMN_X, sides["left"]), (right_x, sides["right"])):
        from_left = card_x == COLUMN_X
        for lane, card in enumerate(group):
            boxes.append(card_html(card, card_x, labels))
            if from_left:
                wire = connector(card, lane, len(group), COLUMN_X + CARD_W, screen_x + card["x"], True)
            else:
                wire = connector(card, lane, len(group), right_x, screen_x + card["x"] + card["w"], False)
            wires.append(wire)

    mid = "center"
    pin = "left:0;top:0;"
    style = (
        f"html,body{{margin:0;padding:0;background:{BACKGROUND};}}"
        f"body{{width:{width}px;height:{height}px;position:relative;font-family:-apple-system,Helvetica,Arial,sans-serif;}}"
        f".shot{{position:absolute;left:{screen_x}px;top:{TOP}px;width:{fw}px;height:{fh}px;}}"
        f".card{{position:absolute;box-sizing:border-box;width:{CARD_W}px;height:{CARD_H}px;background:#fff;"
        "border:2px solid;border-radius:10px;padding:8px 10px;overflow:hidden;color:#1f1f1f;}}"
        ".head{display:flex;gap:6px;align-items:center;margin-bottom:4px;}"
        ".num{color:#fff;border-radius:10px;min-width:20px;height:20px;font-size:12px;font-weight:700;"
        f"display:inline-flex;align-items:{mid};justify-content:{mid};padding:0 4px;}}"
        ".pill{color:#fff;border-radius:9px;font-size:11px;padding:2px 8px;}"
        ".key{font-family:Menlo,monospace;font-size:11px;color:#444;margin-bottom:4px;white-space:nowrap;"
        "overflow:hidden;text-overflow:ellipsis;}"
        ".line{font-size:12px;line-height:16px;max-height:32px;overflow:hidden;}"
        ".muted{color:#666;font-style:italic;}"
        f"svg{{position:absolute;{pin}}}"
    )
    document = (
        f"{PAGE_OPEN}{META_UTF8}<style>{style}</style></head><body>"
        f'<img class="shot" src="{page["image"]}">'
        f'<svg width="{width}" height="{height}">{"".join(wires)}</svg>'
        f'{"".join(boxes)}</body></html>'
    )
    return document, width, height


def find_chrome():
    for path in [os.environ.get("CHROME_BIN")] + MAC_CHROMES:
        if path and os.path.isfile(path) and os.access(path, os.X_OK):
            return path
    return next(filter(None, map(shutil.which, PATH_CHROMES)), None)


def wait_until_stable(proc, path, limit, pause):
    deadline = time.time() + limit
    previous = -1
    while time.time() < deadline and proc.poll() is None:
        current = os.path.getsize(path) if os.path.exists(path) else -1
        if current > 0 and current == previous:
            break
        previous = current
        time.sleep(pause)
    if proc.poll() is None:
        proc.kill()
        proc.wait()


def chrome_screenshot(chrome, html_text, out_png, width, height, scale):
    target = os.path.abspath(out_png)
    with tempfile.TemporaryDirectory() as scratch:
        page = os.path.join(scratch, "page.html")
        with open(page, "w", encoding="utf-8") as sink:
            sink.write(html_text)
        if os.path.exists(target):
            os.remove(target)
        log_path = os.path.join(scratch, "chrome.log")
        window = f"{int(width)},{int(height)}"
        argv = [
            chrome, "--hide-scrollbars", "--headless", "--disable-gpu",
            f"--window-size={window}", f"--force-device-scale-factor={scale}",
            f"--screenshot={target}", f"--user-data-dir={os.path.join(scratch, 'profile')}",
            "file://" + page,
        ]
        with open(log_path, "w") as log:
            proc = subprocess.Popen(argv, stdout=log, stderr=subprocess.STDOUT)
            wait_until_stable(proc, target, 120, 0.5)
        if os.path.isfile(target) and os.path.getsize(target) > 0:
            return
        with open(log_path, errors="replace") as log:
            raise RuntimeError((log.read() or "no output").strip()[-300:])


def is_number(value):
    try:
        float(value)
    except (TypeError, ValueError):
        return False
    return True


def load_spec_pages(path):
    with open(path, encoding="utf-8") as stream:
        spec = json.load(stream)
    folder = os.path.dirname(os.path.abspath(path))
    pages = spec.get("pages") or []
    for number, page in enumerate(pages, 1):
        frame = page.get("frame") if isinstance(page, dict) else None
        if not (isinstance(frame, dict) and all(is_number(frame.get(side)) for side in ("w", "h"))):
            sys.exit(f"{path}: page {number} needs \"frame\": {{\"w\": <width>, \"h\": <height>}} in px")
        for side in ("w", "h"):
            if not isinstance(frame[side], (int, float)):
                frame[side] = float(frame[side])
        page["image"] = image_source(page.get("image"), folder)
    return spec, pages


def split_nodes(flag, mapping):
    given = [piece for piece in (flag or "").split(",") if piece.strip()]
    return given or list(mapping.get("figmaNodes") or [])


def build_cli():
    cli = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    cli.add_argument("--mapping", required=True, help="the reuse-map JSON")
    cli.add_argument("--spec", help="geometry from MCP or measured by hand: {fileKey, pages:[{nodeId, image, frame, nodes}]}")
    cli.add_argument("--file", help="Figma file key; falls back to figmaFileKey in the mapping")
    cli.add_argument("--nodes", help="node-ids separated by commas, the screen frame first; falls back to figmaNodes")
    cli.add_argument("--scale", default=2, type=int)
    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 that receives the generated files")
    cli.add_argument("--slug", help="stem for the file names; derived from the screen name when omitted")
    cli.add_argument("--ui-lang", choices=sorted(LABELS), default="en", help="language of the card labels")
    return cli


def main():
    cli = build_cli()
    opts = cli.parse_args()
    reject_token_flag(cli, opts, FIGMA_TOKEN_HINT)
    labels = LABELS[opts.ui_lang]

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

    if opts.spec:
        spec, pages = load_spec_pages(opts.spec)
        if not pages:
            sys.exit("--spec lists no pages")
        page = pages[0]
        if len(pages) > 1:
            print(f"# warning: {len(pages) - 1} extra page(s) in --spec ignored; the overlay draws the first page only "
                  "(render-key-shots.py reads every page)", file=sys.stderr)
        file_key = spec.get("fileKey") or mapping.get("figmaFileKey")
    else:
        file_key = opts.file or mapping.get("figmaFileKey")
        nodes = split_nodes(opts.nodes, mapping)
        if not (file_key and nodes):
            cli.error("need a Figma file key and frame node-id (flags or mapping), or --spec")
        token = figma_token()
        if not token:
            sys.exit("No Figma token found (FIGMA_ACCESS_TOKEN, FIGMA_TOKEN or keychain FIGMA_ACCESS_TOKEN). "
                     "Without one, measure the boxes on the screenshot and hand them over with --spec.")
        page = page_from_rest(file_key, canon_id(nodes[0]), opts.scale, token, opts.ca_file)

    cards = build_cards(page, mapping.get("rows") or [], labels)
    document, width, height = build_html(page, cards, labels)
    os.makedirs(opts.target_dir, exist_ok=True)

    def output(suffix):
        return os.path.join(opts.target_dir, f"{stem}.overlay.{suffix}")

    chrome = find_chrome()
    if chrome:
        written = output("png")
        try:
            chrome_screenshot(chrome, document, written, width, height, opts.scale)
        except (RuntimeError, subprocess.SubprocessError) as exc:
            sys.exit(f"Chrome could not take the screenshot: {exc}")
    else:
        written = output("html")
        with open(written, "w", encoding="utf-8") as sink:
            sink.write(document)
        print(f"# Chrome, Chromium and Edge are all missing; saved {written} in place of a PNG", file=sys.stderr)

    with_copy = len([card for card in cards if card["annotated"]])
    manifest = {
        "screen": mapping.get("screen"),
        "file": os.path.basename(written),
        "fileKey": file_key,
        "nodeId": page.get("nodeId"),
        "cardCount": len(cards),
        "annotatedCount": with_copy,
        "width": width,
        "height": height,
        "scale": opts.scale,
    }
    with open(output("manifest.json"), "w", encoding="utf-8") as sink:
        json.dump(manifest, sink, indent=2, ensure_ascii=False)
        sink.write("\n")

    if not cards:
        print("# warning: the overlay is empty; no node matched a row and no node had an annotation", file=sys.stderr)
    print(f"# overlay: {len(cards)} card(s), {with_copy} with CMS copy -> {written}", file=sys.stderr)


if __name__ == "__main__":
    main()