#!/usr/bin/env python3
"""roundtrip.py -- convert the radare2 disassembly (qunexus_v2.2.1.asm) into a
single sdas8051 assembler source, assemble it, and verify the result is
byte-identical to the stock firmware image. Optionally apply the USB-1->CV
patch in source form and verify against the expected patched image.

Design:
  * Relative-jump targets (sjmp/jz/jnz/jc/jnc/jb/jnb/jbc/djnz/cjne) are
    computed from the original instruction bytes (addr + len + signed offset),
    NOT parsed from r2's operand text -- this sidesteps r2's `aav.`/`str.`
    symbol quirks. Each target becomes a label `L_XXXX` defined at its byte.
  * Absolute jumps/calls (ljmp/lcall/ajmp/acall) take a raw `0xXXXX` operand.
    ajmp/acall don't 11-bit-encode numerically in sdas8051, so they mis-encode
    and get pinned to `.db` by the verifier (same length -> no layout drift).
  * Bit operands `0xHH.B` are converted to the raw bit address.
  * A few unambiguous direct-of-register forms (push/pop rN, mov @rI,rM) are
    pre-converted; everything else is fed through and pinned to `.db` if
    sdas8051 errors or emits different bytes.
  * Iterative pinning: assemble, find the first layout drift (an instruction
    whose assembled length != original) or any byte mismatch, pin those
    instructions to `.db <orig bytes>`, reassemble, until clean. Converges
    because `.db` always has the exact original length and bytes.

Usage:
  python3 roundtrip.py            # stock round-trip
  python3 roundtrip.py --patch    # apply USB-1->CV patch in source, verify
"""
import re
import subprocess
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
FWTOOLS = HERE.parent
sys.path.insert(0, str(FWTOOLS))   # fallback: parent firmware-tools/
sys.path.insert(0, str(HERE))      # local copy takes precedence (self-contained)
import patch_usb1_to_cv as P  # noqa: E402

BASE = 0x2400
END = 0xF806  # inclusive last byte of the app image
SDAS = "sdas8051"
ASM = "/tmp/rt.asm"
LST = "/tmp/rt.lst"

# relative jumps: (length, offset_byte_index). djnz is special (2 or 3 bytes).
REL_FIXED = {  # op -> (len, off_idx)
    "sjmp": (2, 1), "jz": (2, 1), "jnz": (2, 1), "jc": (2, 1), "jnc": (2, 1),
    "jb": (3, 2), "jnb": (3, 2), "jbc": (3, 2), "cjne": (3, 2),
}
ABS_OPS = {"ljmp", "lcall", "ajmp", "acall"}

LINE_RE = re.compile(
    r'^\s*[\\|/]?\s*0x([0-9a-fA-F]{8})\s+([0-9a-fA-F]+)\s+(\S.*)$')

# Annotated addresses. Emitted as `;` comment lines before the instruction,
# so the disassembly reads as a narrative. These survive regeneration (they
# live in this script, not in the generated file). Patch-mode comments
# (PATCH_COMMENTS) override COMMENTS at the same address.
COMMENTS = {
    0xa57b: "USB-MIDI router (fcn.0000a57b): route the 4-byte event at @r6:r7 to a per-cable destination ring (called from 0xdf28)",
    0xa57d: "switch to register bank 1 (PSW.RS0 = 1)",
    0xa580: "save event-buffer pointer r6:r7 -> 0x0d3b/0x0d3c",
    0xa588: "if CV-source-locked flag (bit 0x22.2) set, early-out (r7 = 1)",
    0xa590: "read event[0]; extract cable number (high nibble) -> r7",
    0xa59c: "cable 0 (USB-1 / Control Surface) -> dest 0x70",
    0xa5a4: "cable 1 -> dest 0xc3",
    0xa5b0: "cable 2 (USB-3) -> dest 0x82 (CV ring)",
    0xa5ba: "write destination descriptor (r3,r2,r1) -> 0x0d3d/e/f  (fcn.0000551e)",
    0xa5c0: "common routing pass: copy MIDI bytes from the event buffer to the selected destination ring (runs ONCE per router invocation)",
    0xa660: "restore PSW; return (r7 = status)",
    0xdf05: "USB-MIDI event dispatcher (fcn.0000df05)",
    0xdf28: "dispatch call -> router 0xa57b",
}

# Patch-mode-only annotations (the stub + retargeted call site).
# Each comment describes the instruction AT that address.
PATCH_COMMENTS = {
    0x8126: "usb1_cv_stub: mirror cable-0 (USB-1) events into the CV ring (dest 0x82); first run the original router unchanged",
    0x8129: "USB-MIDI 4-byte event buffer at 0x0F9B",
    0x812c: "load event[0] (cable nibble in top 4 bits)",
    0x812d: "isolate cable number (top nibble of event[0])",
    0x812f: "not cable 0 -> done (no mirror needed)",
    0x8131: "reload event[0]",
    0x8132: "retag cable 0 as cable 2 (set bit 0x20)",
    0x8134: "store back",
    0x8135: "DPTR high = 0x0F (buffer page)",
    0x8137: "DPTR low  = 0x9B (buffer addr)",
    0x8139: "route again -> USB-3 / CV ring (dest 0x82)",
    0x813c: "done",
    0xdf28: "dispatch call -> usb1_cv_stub  (was: router 0xa57b)",
}


def signed(b):
    return b - 0x100 if b >= 0x80 else b


def parse_asm(path):
    """Return ordered list of (addr, bytes, op, args_str) for instruction lines."""
    out = []
    for ln in Path(path).read_text(errors="replace").splitlines():
        m = LINE_RE.match(ln)
        if not m:
            continue
        addr = int(m.group(1), 16)
        bs = bytes.fromhex(m.group(2))
        rest = m.group(3).split(";", 1)[0].rstrip()
        if not rest:
            continue
        parts = rest.split(None, 1)
        op = parts[0]
        args = parts[1] if len(parts) > 1 else ""
        out.append((addr, bs, op, args))
    return out


def bit_to_raw(byte_hex, bit):
    byte = int(byte_hex, 16)
    if byte < 0x80:
        return (byte - 0x20) * 8 + bit
    return byte + bit


def convert_bits(s):
    def repl(m):
        return "0x%02X" % bit_to_raw(m.group(1), int(m.group(2)))
    return re.sub(r"0x([0-9a-fA-F]+)\.([0-7])", repl, s)


def normalize_symbols(s):
    s = re.sub(r"fcn\.0000([0-9a-fA-F]+)", r"0x\1", s)
    s = re.sub(r"loc\.0000([0-9a-fA-F]+)", r"0x\1", s)
    s = re.sub(r"aav\.0x0000([0-9a-fA-F]+)", r"0x\1", s)
    s = re.sub(r"aav\.0x([0-9a-fA-F]+)", r"0x\1", s)
    return s


def preconvert(op, args):
    """Unambiguous direct-of-register forms that would otherwise error."""
    m = re.match(r"^(push|pop)\s+r([0-7])$", op + " " + args)
    if m:
        return m.group(1), "0x0" + m.group(2)
    m = re.match(r"^mov\s+@r([01]),\s+r([0-7])$", op + " " + args)
    if m:
        return "mov", "@r%s, 0x0%s" % (m.group(1), m.group(2))
    return op, args


def rel_target(addr, bs, op):
    if op == "djnz":
        if len(bs) == 2:
            off = bs[1]
        else:
            off = bs[2]
        return addr + len(bs) + signed(off)
    ln, oi = REL_FIXED[op]
    return addr + ln + signed(bs[oi])


def make_items(parsed, stock):
    """Build the ordered item list covering 0x2400..END.

    Each item: dict(addr, len, orig (bytes), op, args, is_gap, force_db).
    Relative-jump targets are collected and returned as a label set.
    Instruction interiors that contain a label target are split into per-byte
    .db items so the label can land there.
    """
    # First pass: instruction items + gap items, plus label targets.
    labels = set()
    raw = []  # addr -> (bs, op, args, is_gap)
    r2 = {a: (bs, op, args) for (a, bs, op, args) in parsed}
    addr = BASE
    n = END - BASE + 1
    items = []
    # collect relative targets first (need full parse)
    for (a, bs, op, args) in parsed:
        if op in REL_FIXED or op == "djnz":
            t = rel_target(a, bs, op)
            if BASE <= t <= END:
                labels.add(t)
            else:
                labels.add(t)  # keep; will force-db the jump if unplaceable
    # walk addresses
    i = 0
    sorted_instrs = sorted(parsed, key=lambda x: x[0])
    cur = 0
    addr = BASE
    while addr <= END:
        if cur < len(sorted_instrs) and sorted_instrs[cur][0] == addr:
            a, bs, op, args = sorted_instrs[cur]
            # any label strictly inside (a, a+len)?
            interior = [t for t in labels if a < t < a + len(bs)]
            if interior:
                # split into per-byte .db with labels on target bytes
                for k in range(len(bs)):
                    items.append({"addr": a + k, "len": 1, "orig": bs[k:k + 1],
                                  "op": None, "args": "", "is_gap": True,
                                  "force_db": True, "split": True})
            else:
                items.append({"addr": a, "len": len(bs), "orig": bs,
                              "op": op, "args": args, "is_gap": False,
                              "force_db": False})
            addr += len(bs)
            cur += 1
        else:
            items.append({"addr": addr, "len": 1, "orig": stock[addr - BASE:addr - BASE + 1],
                          "op": None, "args": "", "is_gap": True, "force_db": False})
            addr += 1
    return items, labels


def emit_source(items, labels, pinned, patch=False):
    """Emit sdas8051 source. Returns (text, item_source_line) mapping item
    index -> the source line number that produces its bytes (for error maps).

    `items` already reflects the patch (stub + retarget) when patching -- see
    patch_items() -- so emission is uniform: a label before any item whose
    address is a jump/call target, then either a mnemonic or `.db`."""
    lines = []
    lines.append(".area CODE (ABS)")
    lines.append(".org 0x%04x" % BASE)
    item_line = {}
    cmt = dict(COMMENTS)
    if patch:
        cmt.update(PATCH_COMMENTS)
    for idx, it in enumerate(items):
        a = it["addr"]
        if a in cmt:
            lines.append("; " + cmt[a])
        if a in labels:
            lines.append("L_%04X:" % a)
        if it["is_gap"] or it["force_db"] or a in pinned:
            if it["is_gap"] and not it.get("force_db"):
                lines.append(".db 0x%02X" % it["orig"][0])
            else:
                lines.append(".db " + ", ".join("0x%02X" % b for b in it["orig"]))
            item_line[idx] = len(lines)
        else:
            op, args = preconvert(it["op"], it["args"])
            if op in REL_FIXED or op == "djnz":
                t = rel_target(a, it["orig"], op)
                arglist = [x.strip() for x in args.split(",")] if args else []
                if arglist:
                    arglist[-1] = "L_%04X" % t
                arglist = [convert_bits(x) for x in arglist]
                args = ", ".join(arglist)
            else:
                args = convert_bits(args)
            args = normalize_symbols(args)
            lines.append("%s %s" % (op, args) if args else op)
            item_line[idx] = len(lines)
    return "\n".join(lines) + "\n", item_line


# The USB-1->CV stub as (addr, op, args, bytes) items, written into the 0xFF
# padding at 0x8126-0x813C. The jnz target (0x813C, the ret) and the stub
# entry (0x8126, called by the retargeted lcall) become labels.
STUB_ITEMS = [
    (0x8126, "lcall", "0xa57b", b"\x12\xa5\x7b"),
    (0x8129, "mov", "dptr, #0x0f9b", b"\x90\x0f\x9b"),
    (0x812C, "movx", "a, @dptr", b"\xe0"),
    (0x812D, "anl", "a, #0xf0", b"\x54\xf0"),
    (0x812F, "jnz", "L_813C", b"\x70\x0b"),
    (0x8131, "movx", "a, @dptr", b"\xe0"),
    (0x8132, "orl", "a, #0x20", b"\x44\x20"),
    (0x8134, "movx", "@dptr, a", b"\xf0"),
    (0x8135, "mov", "r6, #0x0f", b"\x7e\x0f"),
    (0x8137, "mov", "r7, #0x9b", b"\x7f\x9b"),
    (0x8139, "lcall", "0xa57b", b"\x12\xa5\x7b"),
    (0x813C, "ret", "", b"\x22"),
]


def patch_items(items, labels):
    """Return (items, labels) with the USB-1->CV patch applied: the stock
    0xFF items at 0x8126-0x813C replaced by the stub, and the dispatch call
    at 0xDF28 retargeted to L_8126."""
    stub_lo, stub_hi = P.STUB_ADDR, P.STUB_ADDR + len(P.STUB) - 1
    out = []
    for it in items:
        a = it["addr"]
        if a == stub_lo:
            for (saddr, sop, sargs, sbs) in STUB_ITEMS:
                out.append({"addr": saddr, "len": len(sbs), "orig": sbs,
                            "op": sop, "args": sargs, "is_gap": False,
                            "force_db": False})
            continue
        if stub_lo < a <= stub_hi:
            continue  # covered by the stub block
        if a == P.CALL_SITE:
            out.append({"addr": a, "len": len(P.NEW_CALL), "orig": bytes(P.NEW_CALL),
                        "op": "lcall", "args": "L_%04X" % P.STUB_ADDR,
                        "is_gap": False, "force_db": False})
            continue
        out.append(it)
    labels = set(labels)
    labels.add(P.STUB_ADDR)        # stub entry (called by retarget)
    labels.add(stub_hi)            # jnz target (the ret)
    return out, labels


def parse_lst(path):
    """Return ordered list of (addr, bytes) for byte-producing lines.

    sdas8051 listing: 6-hex addr, single space, hex byte pairs (single-space
    separated), then a 2+ space column gap before `[cycles]` (instruction
    lines) or the decimal line number (`.db` lines). A lazy byte field
    followed by the 2+ space gap avoids grabbing digits from the line number.
    """
    out = []
    rx = re.compile(
        r"^\s+([0-9a-fA-F]{6})\s+([0-9a-fA-F]{2}(?: [0-9a-fA-F]{2})*?)\s{2,}")
    for ln in Path(path).read_text(errors="replace").splitlines():
        m = rx.match(ln)
        if not m:
            continue
        addr = int(m.group(1), 16)
        bs = bytes.fromhex(m.group(2).replace(" ", ""))
        out.append((addr, bs))
    return out


def assemble():
    Path(LST).unlink(missing_ok=True)
    proc = subprocess.run([SDAS, "-l", ASM], capture_output=True, text=True)
    return proc.returncode, proc.stderr


def run(items, labels, patch=False):
    pinned = set()
    # Patch items (stub + retarget) are known to assemble correctly (verified
    # in isolation). They must be exempt from pinning: during early iterations
    # length-mismatch items desync the listing, and these late/inserted items
    # get falsely flagged as mismatches. Left unpinned, they assemble correctly
    # once the real culprits are pinned.
    no_pin = set()
    if patch:
        no_pin.add(P.CALL_SITE)
        no_pin.update(a for (a, _op, _ar, _bs) in STUB_ITEMS)
    for iteration in range(200):
        text, item_line = emit_source(items, labels, pinned, patch=patch)
        Path(ASM).write_text(text)
        rc, stderr = assemble()
        if not Path(LST).exists():
            # no listing -> pin items on error lines
            for ml in re.finditer(r":(\d+): Error", stderr):
                ln = int(ml.group(1))
                for idx, sl in item_line.items():
                    if sl == ln and items[idx]["addr"] not in no_pin:
                        pinned.add(items[idx]["addr"])
            continue
        lst = parse_lst(LST)
        # map error lines to items too
        err_items = set()
        for ml in re.finditer(r":(\d+): Error", stderr):
            ln = int(ml.group(1))
            for idx, sl in item_line.items():
                if sl == ln:
                    err_items.add(idx)
        pins_this = set()
        if err_items:
            # errors desync the listing order; pin them and reassemble
            for idx in err_items:
                if items[idx]["addr"] not in no_pin:
                    pins_this.add(items[idx]["addr"])
        else:
            # no errors -> listing lines map 1:1 to items by emission order.
            # A length mismatch (sdas8051 emitted a different opcode length)
            # would desync the whole tail, so detect it per-item by comparing
            # the assembled length AND bytes -- this pins the actual culprit,
            # not the victim of the drift it caused.
            li = 0
            for idx, it in enumerate(items):
                if li >= len(lst):
                    if it["addr"] not in no_pin:
                        pins_this.add(it["addr"])  # missing tail
                    li += 1
                    continue
                laddr, lbytes = lst[li]
                if len(lbytes) != it["len"] or lbytes != it["orig"]:
                    if it["addr"] not in no_pin:
                        pins_this.add(it["addr"])
                li += 1
        if not pins_this:
            return verify(items, lst, patch), iteration, text
        if iteration < 10 or iteration % 20 == 0:
            print("  iter %d: pins_this=%d err=%d pinned_total=%d" %
                  (iteration, len(pins_this), len(err_items), len(pinned)))
        pinned |= pins_this
    return False, iteration, text


def verify(items, lst, patch):
    # build flat image from lst
    img = bytearray(END - BASE + 1)
    for addr, bs in lst:
        off = addr - BASE
        img[off:off + len(bs)] = bs
    if patch:
        expected = bytearray((HERE / "QuNexus_Firmware_v2.2.1.bin").read_bytes())
        expected[P.STUB_ADDR - BASE:P.STUB_ADDR - BASE + len(P.STUB)] = bytes(P.STUB)
        expected[P.CALL_SITE - BASE:P.CALL_SITE - BASE + len(P.NEW_CALL)] = bytes(P.NEW_CALL)
        expected = bytes(expected)
    else:
        expected = (HERE / "QuNexus_Firmware_v2.2.1.bin").read_bytes()
    ok = bytes(img) == expected
    if not ok:
        # find first diff
        for i in range(min(len(img), len(expected))):
            if img[i] != expected[i]:
                print("  first diff at 0x%04X: got %02X want %02X" %
                      (BASE + i, img[i], expected[i]))
                break
    return ok


def main():
    patch = "--patch" in sys.argv
    out = HERE / "firmware.asm"
    stock = bytearray((HERE / "QuNexus_Firmware_v2.2.1.bin").read_bytes())
    parsed = parse_asm(HERE / "qunexus_v2.2.1.asm")
    items, labels = make_items(parsed, stock)
    if patch:
        items, labels = patch_items(items, labels)
    print("items: %d, labels: %d, instructions: %d" %
          (len(items), len(labels), len(parsed)))
    ok, iters, text = run(items, labels, patch=patch)
    mode = "PATCHED" if patch else "stock"
    if ok:
        out.write_text(text)
        print("PASS: %s image re-assembled byte-identical in %d iterations." % (mode, iters + 1))
        print("wrote %s (%d lines)" % (out.name, text.count("\n")))
        return 0
    print("FAIL: %s image did not match after %d iterations." % (mode, iters + 1))
    return 1


if __name__ == "__main__":
    sys.exit(main())