"""
apply_glue_generic.py - derive stimulus glue for every component of a generated twin
from the twin's own .vsi.cmd, and write it into the skeletons by user-region name.

    python apply_glue_generic.py <twin.vsi.cmd> <twin dir>      e.g. wildfire_grouped.vsi.cmd workspace/WildfireDT

Roles per signal (see generic_stim.h): PRODUCE (source of a `connect signals`, or
any frame signal of the FIRST component defining that frame id), CONSUME (the
destination of a `connect signals`, or the same-position frame signal of every
other component on that id), TRACE (everything else). Every producing BUS (the
generator's "# bus <name> [<protocol>]" group a connect or frame belongs to) gets its
own period from PERIODS in file order, baked into the produce() call, so buses tick
at visibly different rates. Writes <twin dir>/glue_wiring.json (roles, frame owners,
muted CAN consumers, the bus/period table) so check_traces.py can verify a run, and
copies generic_stim.h into <twin dir>/src.
Re-run after any `generate -overwrite`; rebuild with `mingw32-make compile build`.
Exit 1 if any component or region is missing, naming it.
"""

import json
import os
import re
import shutil
import sys
from collections import OrderedDict

HERE = os.path.dirname(os.path.abspath(__file__))
# C++ skeletons mark regions with `//`, Python skeletons with `#`; the region names and
# their placement are the same in both, so one patcher serves both languages
START = re.compile(r"^(\s*)(?://|#) Start of user custom code region\. Please apply edits only within these regions:\s+(.*?)\s*$")
END = re.compile(r"^\s*(?://|#) End of user custom code region\.")


def read_wiring(cmd_path):
    comps = OrderedDict()            # name -> {"signals": [(name, type, dir)], "frames": {id: [sig...]}, "frame_bus": {id: bus}}
    edges = []                       # (src comp, src sig, dst comp, dst sig, bus)
    buses = OrderedDict()            # bus name -> protocol, in file order (the generator's "# bus" comments)
    bus = None                       # the bus the following connectivity lines belong to
    for line in open(cmd_path, encoding="utf-8"):
        line = line.strip()
        m = re.match(r"# bus (\S+) \[(\S+)\]", line)
        if m:
            bus = m.group(1); buses.setdefault(bus, m.group(2))
            continue
        m = re.match(r"add component -type (\S+) -name (\S+)", line)
        if m:
            comps.setdefault(m.group(2), {"type": m.group(1), "signals": [], "frames": OrderedDict(), "frame_bus": {}})
            continue
        m = re.match(r"define componentSignals -componentName (\S+) -signals \[(.*)\]", line)
        if m:
            for s in m.group(2).split(","):
                parts = s.split(":")
                if len(parts) == 3:
                    comps[m.group(1)]["signals"].append(tuple(parts))
            continue
        m = re.match(r"connect signals -sourceSignal (\S+)\.(\S+) -destSignal (\S+)\.(\S+)", line)
        if m:
            edges.append((m.group(1), m.group(2), m.group(3), m.group(4), bus or "unnamed"))
            continue
        m = re.match(r"define frame -componentName (\S+) -portName (\S+) -id (\S+) -signals \[(.*?)\]", line)
        if m:
            sigs = [s.split(":")[0] for s in m.group(4).split(",") if s]
            comps[m.group(1)]["frames"].setdefault(m.group(3), []).extend(sigs)
            comps[m.group(1)]["frame_bus"].setdefault(m.group(3), bus or "frame " + m.group(3))
    return comps, edges, buses


# Distinct periods (steps per counter increment), handed out to producing buses in
# file order so neighbouring buses in the picture tick at visibly different rates.
# With a 100-step run the slowest still changes 8 times, above the checker's floor.
PERIODS = [1, 2, 3, 4, 5, 6, 8, 10, 12]


def plan(comps, edges, buses=None):
    """-> ({comp: {"produce": [(sig, k, period)], "consume": [(sig, producer comp, producer sig)],
                   "trace": [sigs], "silent_can": bool}}, frame_owner, bus_table)
    bus_table: {bus: {"protocol", "period", "signals": [[comp, sig], ...]}} for every bus
    that produces something; the period is what the generated produce() call carries."""
    roles = OrderedDict((c, {"produce": [], "consume": [], "trace": [s[0] for s in d["signals"]]}) for c, d in comps.items())
    kcount = {c: 0 for c in comps}
    bus_table = OrderedDict()

    def period_of(bus, protocol):
        if bus not in bus_table:
            bus_table[bus] = {"protocol": protocol, "period": PERIODS[len(bus_table) % len(PERIODS)], "signals": []}
        return bus_table[bus]["period"]

    def produce(c, sig, bus, protocol):
        if sig not in [p[0] for p in roles[c]["produce"]]:
            roles[c]["produce"].append((sig, kcount[c], period_of(bus, protocol)))
            bus_table[bus]["signals"].append([c, sig])
            kcount[c] += 1

    for sc, ss, dc, ds, bus in edges:
        produce(sc, ss, bus, (buses or {}).get(bus, "?"))
        roles[dc]["consume"].append((ds, sc, ss))
    frame_owner = OrderedDict()     # frame id -> (comp, sigs) of the first definer
    for c, d in comps.items():
        for fid, sigs in d["frames"].items():
            if fid not in frame_owner:
                frame_owner[fid] = (c, sigs)
                fbus = d["frame_bus"].get(fid, "frame " + fid)
                for s in sigs:
                    produce(c, s, fbus, (buses or {}).get(fbus, "CAN"))
            else:
                oc, osigs = frame_owner[fid]
                for i, s in enumerate(sigs):
                    if i < len(osigs):
                        roles[c]["consume"].append((s, oc, osigs[i]))
    # a produced signal is not also consumed on the same component
    for c in roles:
        prod = {p[0] for p in roles[c]["produce"]}
        roles[c]["consume"] = [t for t in roles[c]["consume"] if t[0] not in prod]
    # A CAN node that defines a frame it does not own must not transmit. The generated
    # stub sends its frame unconditionally every step, so a consumer would put a
    # zero-payload frame with the SAME id on the bus and overwrite the producer's
    # values in every other consumer. Measured 2026-09-08 on WildfireDT: 21 consumers
    # of frame 0x100 all traced 0 for all 7 fields while the owner sent 96+step.
    owners = {c for c, _ in frame_owner.values()}
    for c, d in comps.items():
        roles[c]["silent_can"] = bool(d["frames"]) and c not in owners
    return roles, frame_owner, bus_table


def glue_lines(comp, r):
    inc = ['#include "../generic_stim.h"']
    cols = list(r["trace"]) + ["{}_changes".format(s) for s, _, _ in r["consume"]]
    glob = inc + [
        "static long g_step = 0;",
        'static twin::Trace g_trace("{}", {{{}}});'.format(comp, ", ".join('"{}"'.format(c) for c in cols)),
    ]
    if r["consume"]:
        glob.append("static twin::Consumer g_cons[{}];".format(len(r["consume"])))
    before = ["mySignals.{} = twin::produce(g_step, {}, {});".format(s, k, p) for s, k, p in r["produce"]]
    after = ["g_cons[{}].observe(mySignals.{});".format(i, s) for i, (s, _, _) in enumerate(r["consume"])]
    if r.get("silent_can"):
        # The stub's canGateway->sendCanPacket() sits outside every user region, next to
        # the generic-payload sends of the same component, so it cannot be wrapped or
        # skipped: an if(false) block across the regions muted those sends too (measured
        # 2026-09-08, 13 links flat), and an early return would skip the end-of-step
        # handshake and deadlock mainThread. A function-like macro in the globals region
        # rewrites the one call site (canGateway->sendCanPacket()) into a harmless
        # canGateway->setCanId(256) and touches nothing else in the file.
        glob.append("#define sendCanPacket() setCanId(256) /* muted: consumer of a frame it does not own never transmits */")
    vals = ["(long)mySignals.{}".format(s) for s in r["trace"]] + ["(long)g_cons[{}].changes".format(i) for i in range(len(r["consume"]))]
    after.append("g_trace.row(g_step, convert.timeInNs(), {{{}}}); g_step++;".format(", ".join(vals)))
    # Always own all three regions, even when "before" is empty: a region the script
    # does not rewrite keeps whatever an earlier glue version left there (2026-09-08:
    # 19 stale if(false) lines survived a re-apply because the region was skipped).
    regions = {"Global Variables & Definitions": glob, "Before sending the packet": before,
               "After sending the packet": after}
    return regions


def glue_lines_py(comp, r):
    """the Python skeleton's regions carry the same names: "Global Variables &
    Definitions" is module level (before the component class), "Before/After sending
    the packet" sit inside executeMainThreadStep at two tabs. State lives in one
    module-level Glue object so the step can advance without a `global` statement."""
    cols = list(r["trace"]) + ["{}_changes".format(s) for s, _, _ in r["consume"]]
    glob = [
        "import os as _os, sys as _sys",
        "_sys.path.insert(0, _os.path.join(_os.path.dirname(_os.path.abspath(__file__)), '..'))",
        "import generic_stim as twin",
        'g_glue = twin.Glue("{}", [{}], {})'.format(comp, ", ".join('"{}"'.format(c) for c in cols), len(r["consume"])),
    ]
    if r.get("silent_can"):
        # same mute as the C++ macro: the stub calls vsiCanPythonGateway.sendCanPacket()
        # through this module-level name, so rebinding the name to a pass-through
        # wrapper that swallows that one method silences the transmit and nothing else
        glob.append('vsiCanPythonGateway = twin.MutedSend(vsiCanPythonGateway, "sendCanPacket")  # muted: consumer of a frame it does not own')
    before = ["self.mySignals.{} = twin.produce(g_glue.step, {}, {})".format(s, k, p) for s, k, p in r["produce"]]
    after = ["g_glue.cons[{}].observe(self.mySignals.{})".format(i, s) for i, (s, _, _) in enumerate(r["consume"])]
    vals = ["self.mySignals.{}".format(s) for s in r["trace"]] + ["g_glue.cons[{}].changes".format(i) for i in range(len(r["consume"]))]
    after.append("g_glue.trace.row(g_glue.step, vsiCommonPythonApi.getSimulationTimeInNs(), [{}])".format(", ".join(vals)))
    after.append("g_glue.step += 1")
    return {"Global Variables & Definitions": glob, "Before sending the packet": before,
            "After sending the packet": after}


def patch(path, regions):
    with open(path, encoding="utf-8") as fh:
        lines = fh.read().split("\n")
    out, i, seen = [], 0, set()
    while i < len(lines):
        m = START.match(lines[i])
        if not m or m.group(2) not in regions or m.group(2) in seen:
            out.append(lines[i]); i += 1; continue
        indent, name = m.group(1), m.group(2)
        comment = "#" if lines[i].lstrip().startswith("#") else "//"
        out.append(lines[i]); i += 1
        while i < len(lines) and not END.match(lines[i]):
            i += 1
        if i >= len(lines):
            return "region {!r} has no end marker".format(name)
        out.append(indent + comment + " glue from twin_glue/apply_glue_generic.py")
        for g in regions[name]:
            out.append(indent + g)
        out.append(lines[i]); seen.add(name); i += 1
    missing = set(regions) - seen
    if missing:
        return "regions not found: {}".format(sorted(missing))
    with open(path, "w", encoding="utf-8", newline="\n") as fh:
        fh.write("\n".join(out))
    return None


def main():
    if len(sys.argv) != 3:
        sys.exit(__doc__)
    cmd, twin = sys.argv[1], sys.argv[2]
    src = os.path.join(twin, "src")
    if not os.path.isdir(src):
        sys.exit("no src dir under {}".format(twin))
    comps, edges, buses = read_wiring(cmd)
    roles, frame_owner, bus_table = plan(comps, edges, buses)
    shutil.copy(os.path.join(HERE, "generic_stim.h"), os.path.join(src, "generic_stim.h"))
    shutil.copy(os.path.join(HERE, "generic_stim.py"), os.path.join(src, "generic_stim.py"))
    lang = {c: ("Python" if d.get("type") == "Python" else "C++") for c, d in comps.items()}
    wiring = {"cmd": os.path.abspath(cmd), "frame_owners": {fid: c for fid, (c, _) in frame_owner.items()},
              "silent_can": [c for c, r in roles.items() if r.get("silent_can")],
              "buses": bus_table,
              "python_components": [c for c in roles if lang[c] == "Python"],
              "components": {c: {"lang": lang[c], "produce": r["produce"], "consume": r["consume"], "trace": r["trace"]} for c, r in roles.items()}}
    with open(os.path.join(twin, "glue_wiring.json"), "w", encoding="utf-8") as fh:
        json.dump(wiring, fh, indent=1)
    problems, n_prod, n_cons = [], 0, 0
    for comp, r in roles.items():
        py = lang[comp] == "Python"
        path = os.path.join(src, comp, comp + (".py" if py else ".cxx"))
        if not os.path.isfile(path):
            problems.append("{}: {} not found".format(comp, path)); continue
        err = patch(path, glue_lines_py(comp, r) if py else glue_lines(comp, r))
        if err:
            problems.append("{}: {}".format(comp, err))
        elif py:
            # vsiBuild 2026.2 can emit a Python skeleton that is not valid Python: a
            # component with an Ethernet port whose sockets carry no signals got an
            # empty `def establishTcpUdpConnection(self):` (measured 2026-09-08,
            # Gimbal_Video_Aggregator, kept under runs/evidence/). Catch it here, where
            # the file name is known, rather than as a client that never connects.
            import py_compile
            try:
                py_compile.compile(path, doraise=True)
            except py_compile.PyCompileError as e:
                problems.append("{}: generated Python does not compile: {}".format(comp, str(e).strip().splitlines()[-1]))
        n_prod += len(r["produce"]); n_cons += len(r["consume"])
    print("{} components ({} Python): {} produced signal(s), {} consumed signal(s), {} frame id(s) owned by {}; {} CAN consumer(s) muted".format(
        len(roles), len(wiring["python_components"]), n_prod, n_cons, len(frame_owner), ", ".join(c for c, _ in frame_owner.values()), len(wiring["silent_can"])))
    print("bus periods: " + ", ".join("{} [{}] = {}".format(b, t["protocol"], t["period"]) for b, t in bus_table.items()))
    print("wrote", os.path.join(twin, "glue_wiring.json"))
    if problems:
        print("PROBLEMS:")
        for p in problems:
            print("  " + p)
        sys.exit(1)


if __name__ == "__main__":
    main()