#!/usr/bin/env python3
"""Agentic draw.io build engine.

extract_plan(path) : full .drawio  -> compact build-plan (style table + nodes +
                     containers + edges, real geometry & SAP styles, base64
                     image payloads stripped).
build_drawio(plan) : build-plan -> complete, valid .drawio XML (the executor).

The model is trained to emit the build-plan; the engine renders the full file.
Round-tripping every plan guarantees the target is reconstructible & valid.
"""
import base64, zlib, urllib.parse, html, re, json
import xml.etree.ElementTree as ET
from xml.sax.saxutils import quoteattr

IMAGE_DATA_RE = re.compile(r"image=data:[^;,]+;base64,[^;]+;?")
IMG_RE = re.compile(r"image=(data:[^;]+)")  # SAP icons: data:image/png,<b64>
WS_RE = re.compile(r"\s+")

def signature(style):
    """Group styles by icon identity (image payload) else by full string."""
    if not style:
        return "_plain_"
    m = IMG_RE.search(style)
    return ("img:" + m.group(1)) if m else ("sty:" + style.strip())

def _inflate(text):
    text = (text or "").strip()
    if text.startswith("<"):
        return text
    try:
        return urllib.parse.unquote(zlib.decompress(base64.b64decode(text), -15).decode("utf-8"))
    except Exception:
        return None

def _label(val):
    if not val:
        return ""
    return WS_RE.sub(" ", re.sub(r"<[^>]+>", " ", html.unescape(val))).strip()

def _clean_style(style):
    """Strip heavy base64 image payloads but keep SAP shapes/colors/fonts."""
    if not style:
        return "rounded=0;whiteSpace=wrap;html=1;"
    s = IMAGE_DATA_RE.sub("", style)
    s = re.sub(r";{2,}", ";", s).strip(";")
    return (s + ";") if s else "rounded=0;whiteSpace=wrap;html=1;"

def _models(path):
    data = open(path, "r", encoding="utf-8", errors="ignore").read()
    try:
        tree = ET.fromstring(data)
    except ET.ParseError:
        return
    for di, diagram in enumerate(tree.iter("diagram")):
        gm = diagram.find("mxGraphModel")
        if gm is None:
            inf = _inflate(diagram.text)
            if not inf:
                continue
            try:
                gm = ET.fromstring(inf)
            except ET.ParseError:
                continue
        yield diagram.get("name", f"Page-{di+1}"), gm

def extract_plan(path, sig2key, max_pages=1):
    """Return list of build-plans (one per page). Styles -> library keys."""
    plans = []
    for name, gm in _models(path):
        root = gm.find("root")
        if root is None:
            continue
        def sid(style):
            return sig2key.get(signature(style), "box")
        nodes, edges, idmap = [], [], {}
        order = []
        for cell in root.iter("mxCell"):
            cid = cell.get("id")
            if cid in (None, "0", "1"):
                continue
            order.append(cell)
        # assign compact ids
        for i, cell in enumerate(order):
            idmap[cell.get("id")] = i + 2
        for cell in order:
            cid = idmap[cell.get("id")]
            if cell.get("vertex") == "1":
                geo = cell.find("mxGeometry")
                if geo is None:
                    continue
                try:
                    x = int(float(geo.get("x", 0) or 0)); y = int(float(geo.get("y", 0) or 0))
                    w = int(float(geo.get("width", 80) or 80)); h = int(float(geo.get("height", 60) or 60))
                except ValueError:
                    x = y = 0; w, h = 80, 60
                parent = idmap.get(cell.get("parent"), 1)
                n = {"i": cid, "l": _label(cell.get("value")), "k": sid(cell.get("style", "")),
                     "x": x, "y": y, "w": w, "h": h}
                if parent != 1:
                    n["p"] = parent
                nodes.append(n)
            elif cell.get("edge") == "1":
                s = idmap.get(cell.get("source")); t = idmap.get(cell.get("target"))
                if s is None or t is None:
                    continue
                e = {"a": s, "b": t, "k": sid(cell.get("style", ""))}
                lab = _label(cell.get("value"))
                if lab:
                    e["l"] = lab
                edges.append(e)
        if len(nodes) < 2:
            continue
        plans.append({"page": name, "nodes": nodes, "edges": edges})
    return plans[:max_pages] if max_pages else plans

def build_drawio(plan, lib):
    """Executor: build-plan + style library -> full valid .drawio XML string."""
    cells = ['<mxCell id="0" />', '<mxCell id="1" parent="0" />']
    def st(k):
        return lib.get(k, "rounded=0;whiteSpace=wrap;html=1;")
    for n in plan.get("nodes", []):
        parent = n.get("p", 1)
        cells.append(
            f'<mxCell id="{n["i"]}" value={quoteattr(str(n.get("l","")))} '
            f'style={quoteattr(st(n.get("s",0)))} vertex="1" parent="{parent}">'
            f'<mxGeometry x="{n.get("x",0)}" y="{n.get("y",0)}" '
            f'width="{n.get("w",80)}" height="{n.get("h",60)}" as="geometry" /></mxCell>')
    eid = 100000
    for e in plan.get("edges", []):
        lab = e.get("l", "")
        vattr = f' value={quoteattr(str(lab))}' if lab else ''
        cells.append(
            f'<mxCell id="{eid}"{vattr} style={quoteattr(st(e.get("s",0)))} '
            f'edge="1" parent="1" source="{e["a"]}" target="{e["b"]}">'
            f'<mxGeometry relative="1" as="geometry" /></mxCell>')
        eid += 1
    name = plan.get("page", "Architecture")
    return (f'<mxfile host="archgen"><diagram name={quoteattr(name)}>'
            f'<mxGraphModel dx="1200" dy="800" grid="0" gridSize="10" guides="1" '
            f'tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" '
            f'pageWidth="1600" pageHeight="1100" math="0" shadow="0"><root>'
            + "".join(cells) +
            '</root></mxGraphModel></diagram></mxfile>')

def plan_to_text(plan):
    """Compact one-line-per-record serialization the model emits."""
    return json.dumps(plan, ensure_ascii=False, separators=(",", ":"))

def validate(xml):
    """True if well-formed and structurally a real diagram."""
    try:
        t = ET.fromstring(xml)
    except ET.ParseError:
        return False, "parse-error"
    cells = list(t.iter("mxCell"))
    verts = [c for c in cells if c.get("vertex") == "1"]
    return (len(verts) >= 2), f"{len(verts)} nodes, {len([c for c in cells if c.get('edge')=='1'])} edges"

if __name__ == "__main__":
    import glob, os
    lib = json.load(open("dataset/style_library.json"))
    sig2key = json.load(open("dataset/sig2key.json"))
    toks = []
    for f in sorted(glob.glob("data_refarch/**/*.drawio", recursive=True)):
        if "/archive/" in f:
            continue
        for plan in extract_plan(f, sig2key, max_pages=2):
            txt = plan_to_text(plan)
            xml = build_drawio(plan, lib)
            ok, info = validate(xml)
            t = len(txt) // 3
            toks.append(t)
            print(f"{t:6d}tok plan | nodes={len(plan['nodes']):3d} edges={len(plan['edges']):3d} "
                  f"| {'OK ' if ok else 'BAD'} {info:18s} | {os.path.basename(f)[:42]}")
    if toks:
        print(f"\npages={len(toks)} max~tok={max(toks)} "
              f"median~={sorted(toks)[len(toks)//2]} mean~={sum(toks)//len(toks)} "
              f"| >7000tok: {sum(1 for t in toks if t>7000)}")