#!/usr/bin/env python3
"""Semantic spec -> clean, valid SAP .drawio. The model's job shrinks to emitting
typed components + zones + edges; THIS engine does deterministic graphviz layout
and snaps labels to real SAP icon keys. No pixel coords or style keys from the model.

spec = {
  "title": str,
  "zones": [{"id": str, "label": str, "children": [node_id,...]}],
  "nodes": [{"id": str, "label": str, "icon": optional key}],
  "edges": [{"source": id, "target": id, "label": optional}]
}
"""
import json, re, subprocess, difflib
from xml.sax.saxutils import quoteattr

_LIB = json.load(open("dataset/style_library.json"))
_ICON_KEYS = [k for k in _LIB if not k.startswith(("box", "edge"))]

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

def snap_icon(label, given=None):
    """Return a library key for this label (given key wins, else fuzzy match)."""
    if given and given in _LIB:
        return given
    s = slug(label)
    if s in _LIB:
        return s
    m = difflib.get_close_matches(s, _ICON_KEYS, n=1, cutoff=0.72)
    return m[0] if m else None

def _dot_layout(spec):
    """Run graphviz dot, return {node_id: (x,y,w,h)} in pixels, top-left origin."""
    lines = ['digraph G {', 'rankdir=LR; nodesep=0.5; ranksep=0.9;',
             'node [shape=box, fixedsize=true, width=1.1, height=0.7];']
    inzone = set()
    for z in spec.get("zones", []):
        lines.append(f'subgraph cluster_{slug(z["id"])} {{ label={json.dumps(z.get("label",""))};')
        for c in z.get("children", []):
            lines.append(f'"{c}";'); inzone.add(c)
        lines.append('}')
    for n in spec["nodes"]:
        lines.append(f'"{n["id"]}" [label={json.dumps(n["label"])}];')
    for e in spec.get("edges", []):
        lines.append(f'"{e["source"]}" -> "{e["target"]}";')
    lines.append('}')
    dot = "\n".join(lines)
    plain = subprocess.run(["dot", "-Tplain"], input=dot, capture_output=True, text=True).stdout
    scale = 1.0; H = 0.0; pos = {}
    for ln in plain.splitlines():
        p = ln.split()
        if p[0] == "graph":
            scale = float(p[1]); H = float(p[3])
        elif p[0] == "node":
            name = p[1].strip('"')
            x, y, w, h = map(float, p[2:6])
            # graphviz: inches, origin bottom-left -> pixels, top-left
            PX = 72
            pos[name] = (int((x - w/2)*PX), int((H - y - h/2)*PX), int(w*PX), int(h*PX))
    return pos

def _grid_layout(spec):
    """Zone-column layout with SMALLER icons + an OUTER wrapping block. zones =
    columns L->R inside the outer block; nodes stacked. Returns node pos, zone
    rects, and the outer-block rect."""
    NW, NH, GAP, PAD, TITLE, COLGAP, TOP, MARGIN = 96, 54, 30, 18, 30, 46, 64, 30
    zones = spec.get("zones", [])
    zoned = set()
    for z in zones:
        zoned.update(z["children"])
    unzoned = [n["id"] for n in spec["nodes"] if n["id"] not in zoned]
    columns = ([{"id": None, "label": None, "children": unzoned}] if unzoned else []) + \
              [{"id": z["id"], "label": z.get("label", "Zone"), "children": z["children"]} for z in zones]
    pos, zrects = {}, []
    x = TOP + MARGIN
    maxy = TOP
    for col in columns:
        kids = [k for k in col["children"] if any(n["id"] == k for n in spec["nodes"])]
        if not kids:
            continue
        col_w = NW + 2*PAD
        y = TOP + MARGIN + (TITLE if col["id"] else 0) + PAD
        for k in kids:
            pos[k] = (x + PAD, y, NW, NH)
            y += NH + GAP
        if col["id"]:
            zrects.append({"label": col["label"], "rect": [x, TOP + MARGIN, col_w, y - (TOP+MARGIN) - GAP + PAD]})
        maxy = max(maxy, y)
        x += col_w + COLGAP
    outer = [TOP, TOP, x - COLGAP - TOP + MARGIN, maxy - TOP - GAP + PAD + MARGIN]
    return pos, zrects, outer

def build(spec, out_path="spec_output.drawio"):
    pos, zrects, outer = _grid_layout(spec)
    cells = ['<mxCell id="0" />', '<mxCell id="1" parent="0" />']
    # outer surrounding block (behind everything)
    ox, oy, ow, oh = outer
    cells.append(
        f'<mxCell id="outer" value={quoteattr(spec.get("title",""))} '
        f'style="rounded=0;whiteSpace=wrap;html=1;fillColor=#fbfdff;strokeColor=#7a9bbe;'
        f'verticalAlign=top;fontSize=16;fontStyle=1;fontColor=#13334f;dashed=0;" '
        f'vertex="1" parent="1"><mxGeometry x="{ox}" y="{oy}" width="{ow}" height="{oh}" as="geometry" /></mxCell>')
    zid = 1000
    for z in zrects:
        x0, y0, w0, h0 = z["rect"]; zid += 1
        cells.append(
            f'<mxCell id="z{zid}" value={quoteattr(z.get("label") or "")} '
            f'style="rounded=0;whiteSpace=wrap;html=1;fillColor=#eef5fc;strokeColor=#9fc5e8;'
            f'verticalAlign=top;fontSize=13;fontStyle=1;fontColor=#1a3b5c;" vertex="1" parent="1">'
            f'<mxGeometry x="{x0}" y="{y0}" width="{w0}" height="{h0}" as="geometry" /></mxCell>')
    idmap = {}; nid = 1
    for n in spec["nodes"]:
        if n["id"] not in pos:
            continue
        x, y, w, h = pos[n["id"]]
        nid += 1; idmap[n["id"]] = f"n{nid}"
        key = snap_icon(n["label"], n.get("icon"))
        if key:
            style = _LIB[key].rstrip(";") + ";verticalLabelPosition=bottom;verticalAlign=top;"
        else:
            style = "rounded=1;whiteSpace=wrap;html=1;fillColor=#ffffff;strokeColor=#5b7a9c;"
        cells.append(
            f'<mxCell id="n{nid}" value={quoteattr(n["label"])} style={quoteattr(style)} '
            f'vertex="1" parent="1"><mxGeometry x="{x}" y="{y}" width="{w}" height="{h}" '
            f'as="geometry" /></mxCell>')
    # edges
    eid = 5000
    for e in spec.get("edges", []):
        if e["source"] not in idmap or e["target"] not in idmap:
            continue
        eid += 1
        lab = e.get("label", "")
        va = f' value={quoteattr(lab)}' if lab else ''
        cells.append(
            f'<mxCell id="e{eid}"{va} style="edgeStyle=orthogonalEdgeStyle;rounded=0;html=1;'
            f'strokeColor=#5b7a9c;" edge="1" parent="1" source={quoteattr(idmap[e["source"]])} '
            f'target={quoteattr(idmap[e["target"]])}><mxGeometry relative="1" as="geometry" /></mxCell>')
    xml = (f'<mxfile host="archgen"><diagram name={quoteattr(spec.get("title","Architecture"))}>'
           f'<mxGraphModel dx="1200" dy="800" grid="0" guides="1" tooltips="1" connect="1" '
           f'arrows="1" fold="1" page="1" pageScale="1" pageWidth="1600" pageHeight="1100" '
           f'math="0" shadow="0"><root>' + "".join(cells) +
           '</root></mxGraphModel></diagram></mxfile>')
    open(out_path, "w").write(xml)
    return out_path

# ---------------- recursive nested-block layout (hierarchical spec) ----------------
# wide cell (gives the label room) + SMALL icon; row pitch leaves space for the
# bottom-positioned 2-line label so it never overlaps the next icon.
ICON_W, ICON_H, BPAD, BTITLE = 124, 40, 20, 34
COLGAP, ROWGAP, LABEL_H = 34, 26, 34
BGAP = COLGAP  # block-to-block spacing
ROW_PITCH = ICON_H + LABEL_H + ROWGAP

def _measure(block):
    cw, ch = [], []
    for sub in block.get("blocks", []):
        _measure(sub); cw.append(sub["_w"]); ch.append(sub["_h"])
    comps = block.get("components", [])
    if comps:
        ncol = 1 if len(comps) <= 5 else 2
        rows = (len(comps) + ncol - 1) // ncol
        clw = ncol*ICON_W + (ncol-1)*COLGAP
        clh = rows*ROW_PITCH - ROWGAP            # no trailing gap; last row keeps its label
        block["_cluster"] = (ncol, clw, clh); cw.append(clw); ch.append(clh)
    inner_w = (sum(cw) + COLGAP*(len(cw)-1)) if cw else ICON_W
    inner_h = max(ch) if ch else ICON_H
    block["_w"] = inner_w + 2*BPAD
    block["_h"] = inner_h + BTITLE + BPAD

def _place(block, x, y, pos, rects, depth=0):
    rects.append({"label": block.get("label", ""), "rect": [x, y, block["_w"], block["_h"]], "depth": depth})
    cx = x + BPAD; cy = y + BTITLE
    for sub in block.get("blocks", []):
        _place(sub, cx, cy, pos, rects, depth+1); cx += sub["_w"] + COLGAP
    comps = block.get("components", [])
    if comps:
        ncol = block["_cluster"][0]
        for i, comp in enumerate(comps):
            r, c = divmod(i, ncol)
            pos[comp["id"]] = (cx + c*(ICON_W+COLGAP), cy + r*ROW_PITCH, ICON_W, ICON_H)

def compute_layout(spec):
    """Return {component_id: (x, y, w, h)} icon-cell positions (no file written)."""
    return compute_full(spec)[0]

def compute_full(spec):
    """Return (pos {id:(x,y,w,h)}, block_rects [{label,rect,depth}]) — for streaming/render."""
    import copy as _copy
    blocks = _copy.deepcopy(spec.get("blocks", []))
    for b in blocks: _measure(b)
    pos, rects = {}, []; x = 50
    for b in blocks:
        _place(b, x, 60, pos, rects, 1); x += b["_w"] + 50
    return pos, rects

def layout_issues(spec, tol=4):
    """Visual sanity check: flag any two component boxes (icon + label area) that overlap.
    The engine is overlap-free by construction, so this should always pass."""
    pos = compute_layout(spec)
    # each node occupies icon cell + the label strip below it (LABEL_H tall)
    boxes = {cid: (x, y, w, ICON_H + LABEL_H) for cid, (x, y, w, h) in pos.items()}
    ids = list(boxes); out = []
    for i in range(len(ids)):
        ax, ay, aw, ah = boxes[ids[i]]
        for j in range(i + 1, len(ids)):
            bx, by, bw, bh = boxes[ids[j]]
            if ax < bx+bw-tol and bx < ax+aw-tol and ay < by+bh-tol and by < ay+ah-tol:
                out.append(f"layout overlap: {ids[i]} & {ids[j]}")
    return out[:6]

def _fix_img(style):
    # browsers/mxGraph need ;base64 for base64 data URIs (draw.io tolerates the comma form)
    return re.sub(r"data:image/(svg\+xml|png|jpeg),", r"data:image/\1;base64,", style)

def stream_cells(spec):
    """Real draw.io cells (blocks → nodes → edges) in build order, with EXACT styles +
    geometry — for incremental insertion into an mxGraph (pixel-identical to build_hier)."""
    pos, rects = compute_full(spec)
    fills = ["#eef5fc", "#e4eefb", "#dbe8f8"]
    idlabel = {}
    def walk(bs):
        for b in bs:
            for c in b.get("components", []): idlabel[c["id"]] = c["label"]
            walk(b.get("blocks", []))
    walk(spec.get("blocks", []))
    cells = []
    for i, z in enumerate(sorted(rects, key=lambda r: r["depth"])):
        x, y, w, h = z["rect"]; d = z["depth"]
        fc = fills[min(d-1, len(fills)-1)] if d >= 1 else "#ffffff"
        cells.append({"id": f"b{i}", "vertex": True, "value": z.get("label") or "",
                      "x": x, "y": y, "w": w, "h": h,
                      "style": f"rounded=0;whiteSpace=wrap;html=1;fillColor={fc};strokeColor=#9fc5e8;"
                               f"verticalAlign=top;fontSize=12;fontStyle=1;fontColor=#1a3b5c;"})
    nmap = {}
    for cid, (x, y, w, h) in pos.items():
        lbl = idlabel.get(cid, ""); key = snap_icon(lbl)
        # use the EXACT style build_hier uses (draw.io renders the original comma data-URI form)
        style = (_LIB[key].rstrip(";") + ";verticalLabelPosition=bottom;verticalAlign=top;"
                 "imageAlign=center;align=center;whiteSpace=wrap;") if key \
                else "rounded=1;whiteSpace=wrap;html=1;fillColor=#ffffff;strokeColor=#5b7a9c;align=center;"
        nmap[cid] = f"n{cid}"
        cells.append({"id": nmap[cid], "vertex": True, "value": lbl,
                      "x": x, "y": y, "w": w, "h": ICON_H, "style": style})
    for i, e in enumerate(spec.get("connections", [])):
        if e["source"] in nmap and e["target"] in nmap:
            cells.append({"id": f"e{i}", "edge": True, "value": e.get("label", ""),
                          "source": nmap[e["source"]], "target": nmap[e["target"]],
                          "style": "edgeStyle=orthogonalEdgeStyle;rounded=0;html=1;strokeColor=#42607d;endArrow=block;"})
    return cells

def build_hier(spec, out_path="hier_output.drawio"):
    import copy as _copy
    blocks = _copy.deepcopy(spec.get("blocks", []))   # don't pollute caller's spec with _w/_h
    for b in blocks: _measure(b)
    pos, rects = {}, []
    x = 50
    maxh = 0
    placed = []
    for b in blocks:
        _place(b, x, 60, pos, rects, 1); x += b["_w"] + 50; maxh = max(maxh, b["_h"])
    # outer wrapper
    ow = x - 50 - 50 + 2*30; oh = maxh + 2*30
    cells = ['<mxCell id="0" />', '<mxCell id="1" parent="0" />']
    cells.append(f'<mxCell id="outer" value={quoteattr(spec.get("title",""))} '
                 f'style="rounded=0;whiteSpace=wrap;html=1;fillColor=#ffffff;strokeColor=#5b7a9c;'
                 f'verticalAlign=top;fontSize=17;fontStyle=1;fontColor=#13334f;" vertex="1" parent="1">'
                 f'<mxGeometry x="20" y="20" width="{max(ow, x-50)}" height="{oh+60}" as="geometry" /></mxCell>')
    fills = ["#eef5fc", "#e4eefb", "#dbe8f8"]
    for i, z in enumerate(sorted(rects, key=lambda r: r["depth"])):
        x0, y0, w0, h0 = z["rect"]
        fc = fills[min(z["depth"]-1, len(fills)-1)] if z["depth"] >= 1 else "#ffffff"
        cells.append(f'<mxCell id="blk{i}" value={quoteattr(z.get("label") or "")} '
                     f'style="rounded=0;whiteSpace=wrap;html=1;fillColor={fc};strokeColor=#9fc5e8;'
                     f'verticalAlign=top;fontSize=12;fontStyle=1;fontColor=#1a3b5c;" vertex="1" parent="1">'
                     f'<mxGeometry x="{x0}" y="{y0}" width="{w0}" height="{h0}" as="geometry" /></mxCell>')
    idmap = {}; nid = 1
    def walk(b):
        for sub in b.get("blocks", []): walk(sub)
        for comp in b.get("components", []):
            if comp["id"] not in pos: continue
            nonlocal nid; nid += 1; idmap[comp["id"]] = f"n{nid}"
            xx, yy, ww, hh = pos[comp["id"]]
            key = snap_icon(comp["label"], comp.get("icon"))
            # small icon centered in the wide cell; label wraps below within the full width
            style = (_LIB[key].rstrip(";") + ";verticalLabelPosition=bottom;verticalAlign=top;"
                     "imageAlign=center;labelPosition=center;align=center;whiteSpace=wrap;") if key \
                    else "rounded=1;whiteSpace=wrap;html=1;fillColor=#ffffff;strokeColor=#5b7a9c;align=center;"
            cells.append(f'<mxCell id="n{nid}" value={quoteattr(comp["label"])} style={quoteattr(style)} '
                         f'vertex="1" parent="1"><mxGeometry x="{xx}" y="{yy}" width="{ww}" height="{hh}" as="geometry" /></mxCell>')
    for b in blocks: walk(b)
    eid = 5000
    for e in spec.get("connections", []):
        if e["source"] not in idmap or e["target"] not in idmap: continue
        eid += 1; lab = e.get("label", ""); va = f' value={quoteattr(lab)}' if lab else ''
        cells.append(f'<mxCell id="e{eid}"{va} style="edgeStyle=orthogonalEdgeStyle;rounded=0;html=1;'
                     f'strokeColor=#42607d;endArrow=block;" edge="1" parent="1" '
                     f'source={quoteattr(idmap[e["source"]])} target={quoteattr(idmap[e["target"]])}>'
                     f'<mxGeometry relative="1" as="geometry" /></mxCell>')
    xml = (f'<mxfile host="archgen"><diagram name={quoteattr(spec.get("title","Architecture"))}>'
           f'<mxGraphModel dx="1200" dy="800" grid="0" guides="1" tooltips="1" connect="1" arrows="1" '
           f'fold="1" page="1" pageScale="1" pageWidth="1600" pageHeight="1100" math="0" shadow="0"><root>'
           + "".join(cells) + '</root></mxGraphModel></diagram></mxfile>')
    open(out_path, "w").write(xml)
    return out_path

if __name__ == "__main__":
    import sys
    spec = json.load(open(sys.argv[1])) if len(sys.argv) > 1 else json.load(open("demo_spec.json"))
    fn = build_hier if "blocks" in spec else build
    p = fn(spec, sys.argv[2] if len(sys.argv) > 2 else "spec_output.drawio")
    print("wrote", p)