#!/usr/bin/env python3
"""Generate docs/agent/ba-drawio.md from the BA's draw.io source file.

Source: the newest "TNSU_EDU *.drawio" in the repository root (or a path given
on the command line). Each draw.io page is rendered as Markdown:

- Pages whose shapes are connected by arrows become nested lists.
- Unconnected numbered items such as "4.3 ..." are attached to the module
  whose text starts with "4.".
- Unconnected text boxes are attached to the connected shape directly above
  them in the same column (used for the curriculum page).
- Pages built from grouped step boxes (flows) become numbered step lists, with
  side notes attached to the step on the same row.

Re-run after the BA updates the file:

    python3 docs/agent/tools/gen_ba_drawio.py

Uses the standard library only.
"""

import base64
import html
import re
import sys
import urllib.parse
import xml.etree.ElementTree as ET
import zlib
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
OUT = ROOT / "docs" / "agent" / "ba-drawio.md"

COLOUR_NAMES = {
    "#fff2cc": "yellow",
    "#dae8fc": "blue",
    "#d5e8d4": "green",
    "#f8cecc": "pink",
    "#e1d5e7": "purple",
    "#ffe6cc": "orange",
}


def find_source():
    if len(sys.argv) > 1:
        return Path(sys.argv[1])
    files = sorted(ROOT.glob("TNSU_EDU*.drawio"), key=lambda p: p.stat().st_mtime)
    if not files:
        sys.exit("no TNSU_EDU*.drawio file in the repository root")
    return files[-1]


def clean(value):
    v = re.sub(r"<br\s*/?>|</div>|</p>|</li>", "\n", value or "", flags=re.I)
    v = html.unescape(re.sub(r"<[^>]+>", "", v)).replace("\xa0", " ")
    lines = []
    for line in v.split("\n"):
        # "Title- detail" written without a line break: split before the dash
        lines += re.split(r"(?<=\S)(?=-\s)", line)
    lines = [re.sub(r"\s+", " ", line).strip() for line in lines]
    return [line for line in lines if line]


def number_key(lines):
    """Sort key from a leading "8." or "4.3" number; unnumbered items keep diagram order."""
    m = re.match(r"(\d+)(?:\.(\d+))?", lines[0]) if lines else None
    return (int(m.group(1)), int(m.group(2) or 0)) if m else None


def style(cell, key):
    m = re.search(r"(?:^|;)" + key + r"=([^;]*)", cell.get("style") or "")
    return m.group(1) if m else ""


def page_root(diagram):
    model = diagram.find("mxGraphModel")
    if model is not None:
        return model.find("root")
    data = zlib.decompress(base64.b64decode((diagram.text or "").strip()), -15).decode()
    return ET.fromstring(urllib.parse.unquote(data)).find("root")


class Page:
    def __init__(self, diagram):
        self.name = diagram.get("name")
        root = page_root(diagram)
        self.cells = {c.get("id"): c for c in root.iter("mxCell")}
        self.edges = [c for c in self.cells.values() if c.get("edge") == "1"]
        self.verts = {cid: c for cid, c in self.cells.items() if c.get("vertex") == "1"}

    def geo(self, cid):
        """Absolute x, y, w, h (child geometry is relative to its parent group)."""
        c = self.cells[cid]
        g = c.find("mxGeometry")
        x, y, w, h = (float(g.get(k, 0)) for k in ("x", "y", "width", "height")) if g is not None else (0, 0, 0, 0)
        parent = c.get("parent")
        if parent in self.verts:
            px, py, _, _ = self.geo(parent)
            x, y = x + px, y + py
        return x, y, w, h

    def text(self, cid):
        return clean(self.cells[cid].get("value"))

    def kind(self, cid):
        return (self.cells[cid].get("style") or "").split(";")[0]


def colour(cell):
    fill = style(cell, "fillColor").lower()
    return COLOUR_NAMES.get(fill, fill if fill and fill != "none" else "")


def bullet(lines):
    """First line as the item; following lines as its detail."""
    head, *rest = lines
    if rest:
        return head + " (" + "; ".join(r.lstrip("- ").strip() for r in rest) + ")"
    return head


def render_tree(page):
    children = defaultdict(list)
    has_parent = set()
    for e in page.edges:
        s, t = e.get("source"), e.get("target")
        if s in page.verts and t in page.verts and s != t:
            children[s].append(t)
            has_parent.add(t)
    connected = set(children) | has_parent

    # Numbered loose items ("4.3 ...") belong to the module whose text starts with "4."
    by_number = {}
    for cid in connected:
        lines = page.text(cid)
        m = re.match(r"(\d+)\.\s", lines[0]) if lines else None
        if m:
            by_number[m.group(1)] = cid
    extra = defaultdict(list)
    loose = []
    for cid in page.verts:
        if cid in connected or not page.text(cid) or page.kind(cid) in ("group", "swimlane"):
            continue
        lines = page.text(cid)
        m = re.match(r"(\d+)\.\d+", lines[0])
        if m and m.group(1) in by_number:
            extra[by_number[m.group(1)]].append(cid)
        else:
            loose.append(cid)

    # Loose text boxes go under the connected shape directly above them in the same column.
    for cid in loose:
        x, y, _, _ = page.geo(cid)
        best = None
        for other in connected:
            ox, oy, _, oh = page.geo(other)
            if abs(ox - x) <= 40 and oy + oh <= y + 5:
                if best is None or oy > page.geo(best)[1]:
                    best = other
        extra[best].append(cid)

    out = []

    def show(cid, depth, seen):
        c = page.cells[cid]
        tag = f" `{colour(c)}`" if colour(c) else ""
        out.append("  " * depth + f"- {bullet(page.text(cid)) or '(no text)'}{tag}")
        if cid in seen:
            return
        seen.add(cid)
        kids = children[cid]
        if kids and all(number_key(page.text(k)) for k in kids):
            kids = sorted(kids, key=lambda k: number_key(page.text(k)))
        else:
            kids = sorted(kids, key=lambda k: (page.geo(k)[0], page.geo(k)[1]))
        for k in kids:
            show(k, depth + 1, seen)
        for k in sorted(extra.get(cid, []), key=lambda k: number_key(page.text(k)) or (999, page.geo(k)[1])):
            lines = page.text(k)
            tag = f" `{colour(page.cells[k])}`" if colour(page.cells[k]) else ""
            if number_key(lines):  # numbered sub-item: title plus its details
                out.append("  " * (depth + 1) + f"- {bullet(lines)}{tag}")
                continue
            # A loose box may hold several items, each starting with "-".
            items, cur = [], []
            for line in lines:
                if line.startswith("-") and cur:
                    items.append(cur)
                    cur = []
                cur.append(line)
            items.append(cur)
            for it in items:
                out.append("  " * (depth + 1) + f"- {' '.join(l.lstrip('- ').strip() for l in it)}{tag}")

    seen = set()
    roots = [cid for cid in children if cid not in has_parent]
    for r in sorted(roots, key=lambda r: (page.geo(r)[1], page.geo(r)[0])):
        show(r, 0, seen)
    for cid in extra.get(None, []):
        out.append(f"- {bullet(page.text(cid))}")
    return out


def render_flow(page):
    groups = [cid for cid in page.verts if page.kind(cid) == "group"]
    groups.sort(key=lambda g: page.geo(g)[1])
    lanes = [cid for cid in page.verts if page.kind(cid) == "swimlane"]
    notes = [cid for cid in page.verts if page.kind(cid) == "text" and page.text(cid)]
    out = []
    if lanes:
        out.append(f"Swimlane: **{' / '.join(page.text(lanes[0]))}**")
        out.append("")
    for i, g in enumerate(groups, 1):
        members = [cid for cid, c in page.verts.items() if c.get("parent") == g]
        members.sort(key=lambda m: page.geo(m)[1])
        step = " ".join(page.text(members[0])) if members else ""
        module = " ".join(page.text(members[1])) if len(members) > 1 else ""
        gx, gy, gw, gh = page.geo(g)
        out.append(f"{i}. **{step}**" + (f" (sub-module: {module})" if module else ""))
        for n in notes:
            nx, ny, _, _ = page.geo(n)
            if gy - 20 <= ny <= gy + gh + 30 and abs(nx - gx) < 300:
                lines = page.text(n)
                parts = [p for line in lines for p in re.split(r"(?<=\S)-\s", line)]
                for p in parts:
                    out.append(f"   - {p.lstrip('- ').strip()}")
    return out


def main():
    src = find_source()
    pages = [Page(d) for d in ET.parse(src).getroot().findall("diagram")]
    lines = [
        "# BA draw.io source (generated)",
        "",
        f"> Generated by `docs/agent/tools/gen_ba_drawio.py` from `{src.name}`. "
        "Do not edit by hand; re-run the script after the BA updates the file. "
        "For what this means for the project, read [03-modules.md](03-modules.md).",
        "",
        "Colours are the fill colours used in the diagram (their meaning has not been stated by the BA).",
        "",
    ]
    for page in pages:
        lines += [f"## {page.name}", ""]
        is_flow = any(page.kind(c) == "group" for c in page.verts)
        lines += render_flow(page) if is_flow else render_tree(page)
        lines.append("")
    OUT.write_text("\n".join(lines), encoding="utf-8")
    print(f"Wrote {OUT.relative_to(ROOT)} from {src.name} ({len(pages)} pages)")


if __name__ == "__main__":
    main()