#!/usr/bin/env python3
"""Patch QuNexus firmware so USB port 1 MIDI also drives the CV outputs.

Background (all addresses are in the 8051 code space of the application image,
which the bootloader programs from 0x2400 upward):

  0xD3C0  USB endpoint-2 OUT service loop: reads a 4-byte USB-MIDI event from
          FIFO2 into XDATA 0x0F9B, then calls the dispatcher at 0xDF05.
  0xDF05  computes the CIN length, then calls the router at 0xA57B (its only
          caller, via `LCALL 0xA57B` at 0xDF28).
  0xA57B  reads byte 0 of the event, takes the cable number (SWAP A / ANL A,#0Fh
          at 0xA596) and selects the destination ring buffer:
              cable 0 -> 0x0370   USB port 1 "Control Surface"
              cable 1 -> 0x01C3   USB port 2 "Expander"
              cable 2 -> 0x0382   USB port 3, the port feeding the CV engine

Because the CV_Out_*_MIDI_Input_Device preset enum only offers Expander / USB 3
(see CV_Out_Source in qt-qunexus/source/midiio/sysexencdecode.cpp:652), MIDI
arriving on USB 1 can never reach the CV outputs.

This patch redirects the single call site at 0xDF28 to a stub placed in unused
flash at 0xE8F2. The stub runs the original routing first (so USB 1 keeps every
existing behaviour), and then, only for cable 0, rewrites the cable nibble to 2
and routes the same event a second time -- into the USB-3/CV ring. Net effect:
USB 1 events are delivered to both their normal destination and the CV engine.
"""

import argparse
import sys

# ---------------------------------------------------------------- syx container

def sysex_messages(data):
    msgs, i = [], 0
    while True:
        s = data.find(b"\xf0", i)
        if s < 0:
            break
        e = data.find(b"\xf7", s)
        if e < 0:
            break
        msgs.append(data[s:e + 1])
        i = e + 1
    return msgs


def decode_7in8(buf):
    out = bytearray()
    for i in range(0, len(buf) - 7, 8):
        hi = buf[i + 7]
        for j in range(7):
            out.append(buf[i + j] | (0x80 if (hi >> j) & 1 else 0))
    return bytes(out)


def encode_7in8(buf):
    """Inverse of midi_sx_encode_char(); caller must pad buf to a multiple of 7."""
    assert len(buf) % 7 == 0
    out = bytearray()
    for i in range(0, len(buf), 7):
        grp = buf[i:i + 7]
        hi = 0
        for j, c in enumerate(grp):
            out.append(c & 0x7F)
            if c & 0x80:
                hi |= 1 << j
        out.append(hi)
    return bytes(out)


def split_message(msg):
    """-> (prefix bytes through SX_PACKET_START, decoded payload)."""
    body = msg[1:-1]
    i = 6
    while i < len(body) and body[i] == 0x00:
        i += 1
    assert body[i] == 0x01, "SX_PACKET_START not found"
    return msg[:1 + i + 1], decode_7in8(body[i + 1:])


def rebuild_message(prefix, payload):
    pad = (-len(payload)) % 7
    return prefix + encode_7in8(payload + b"\x00" * pad) + b"\xf7"


def crc16(data, crc=0xFFFF):
    """SysExEncDecode::crc_byte (sysexencdecode.cpp:1878), seed 0xFFFF."""
    for ch in data:
        temp = ((crc >> 8) ^ ch) & 0xFFFF
        crc = (crc << 8) & 0xFFFF
        quick = (temp ^ (temp >> 4)) & 0xFFFF
        crc = (crc ^ quick) & 0xFFFF
        quick = (quick << 5) & 0xFFFF
        crc = (crc ^ quick) & 0xFFFF
        quick = (quick << 7) & 0xFFFF
        crc = (crc ^ quick) & 0xFFFF
    return crc


def make_packet(addr, data, rtype=0x00):
    """Build one framed record: 03 LEN 3A LL AAAA TT <data> CC CRChi CRClo."""
    rec = bytes([len(data), (addr >> 8) & 0xFF, addr & 0xFF, rtype]) + data
    rec = b"\x3a" + rec + bytes([(-sum(rec)) & 0xFF])
    body = bytes([0x03, len(rec) + 1]) + rec
    c = crc16(body)
    return body + bytes([c >> 8, c & 0xFF])


def iter_packets(payload):
    """Yield (start, end, addr, rtype, datalen) for each packet in a payload."""
    i = 7
    while i + 1 < len(payload):
        while i < len(payload) and payload[i] == 0x00:
            i += 1
        if i + 1 >= len(payload) or payload[i] != 0x03:
            break
        ln = payload[i + 1]
        end = i + 1 + ln + 2
        addr = (payload[i + 4] << 8) | payload[i + 5]
        yield (i, end, addr, payload[i + 6], ln - 7)
        i = end


def reseal(payload, start, end):
    """Recompute the hex checksum and CRC16 of the packet at [start:end)."""
    ln = payload[start + 1]
    span = payload[start + 3:start + ln]          # LL AAAA TT data
    payload[start + ln] = (-sum(span)) & 0xFF     # CC (Intel-HEX checksum)
    c = crc16(bytes(payload[start:start + 1 + ln]))
    payload[start + 1 + ln] = c >> 8
    payload[start + 2 + ln] = c & 0xFF


# ---------------------------------------------------------------- the patch

# Where the stub goes. This MUST be flash that a stock hex record already
# covers, otherwise the bootloader may never erase/program it and the LCALL
# lands in unprogrammed flash. 0x8126 is 25 bytes of 0xFF linker padding
# between two data tables and is inside a stock record; 0xE8F2 (a 270-byte
# hole covered by NO record) was tried first and bricked MIDI input, because
# the added record was not programmed. Do not use uncovered gaps.
STUB_ADDR = 0x8126
STUB_MAX = 25               # size of the 0xFF run at STUB_ADDR
CALL_SITE = 0xDF28          # LCALL 0xA57B inside the USB-MIDI dispatcher 0xDF05
ROUTER = 0xA57B
EVENT_BUF = 0x0F9B          # XDATA holding the 4-byte USB-MIDI event

STUB = bytes([
    0x12, ROUTER >> 8, ROUTER & 0xFF,        # LCALL 0xA57B      normal routing
    0x90, EVENT_BUF >> 8, EVENT_BUF & 0xFF,  # MOV DPTR,#0x0F9B
    0xE0,                                    # MOVX A,@DPTR      A = byte0
    0x54, 0xF0,                              # ANL A,#0F0h       cable nibble
    0x70, 0x0B,                              # JNZ  done         not cable 0
    0xE0,                                    # MOVX A,@DPTR
    0x44, 0x20,                              # ORL A,#020h       cable 0 -> 2
    0xF0,                                    # MOVX @DPTR,A
    0x7E, EVENT_BUF >> 8,                    # MOV R6,#00Fh
    0x7F, EVENT_BUF & 0xFF,                  # MOV R7,#09Bh
    0x12, ROUTER >> 8, ROUTER & 0xFF,        # LCALL 0xA57B      -> CV ring
    0x22,                                    # done: RET
])
NEW_CALL = bytes([0x12, STUB_ADDR >> 8, STUB_ADDR & 0xFF])
OLD_CALL = bytes([0x12, ROUTER >> 8, ROUTER & 0xFF])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("infile")
    ap.add_argument("outfile")
    args = ap.parse_args()

    data = open(args.infile, "rb").read()
    msgs = [split_message(m) for m in sysex_messages(data)]
    msgs = [(p, bytearray(pl)) for p, pl in msgs]

    # --- 1. retarget the LCALL at 0xDF28 -------------------------------------
    # An instruction can straddle two hex records, so write byte-wise and
    # reseal every record touched.
    index = []          # (addr, dlen, payload, start, end)
    for _prefix, payload in msgs:
        for start, end, addr, rtype, dlen in iter_packets(payload):
            if rtype == 0x00:
                index.append((addr, dlen, payload, start, end))

    def locate(a):
        for addr, dlen, payload, start, end in index:
            if addr <= a < addr + dlen:
                return payload, start, end, start + 7 + (a - addr)
        return None

    touched = {}

    def write(addr, new, expect=None):
        """Write bytes at `addr` into whatever record(s) already cover them."""
        for k, b in enumerate(new):
            loc = locate(addr + k)
            if loc is None:
                raise SystemExit(
                    f"ERROR: 0x{addr + k:04X} is not covered by any hex record. "
                    f"The stub must live in flash a stock record already writes.")
            payload, start, end, off = loc
            if expect is not None and payload[off] != expect[k]:
                raise SystemExit(
                    f"ERROR: expected 0x{expect[k]:02X} at 0x{addr + k:04X}, "
                    f"found 0x{payload[off]:02X}")
            payload[off] = b
            touched[(id(payload), start)] = (payload, start, end)

    if len(STUB) > STUB_MAX:
        raise SystemExit(f"ERROR: stub is {len(STUB)} bytes, only {STUB_MAX} free")

    # 1. the stub, into existing 0xFF padding (verify it really is free first)
    write(STUB_ADDR, STUB, expect=b"\xff" * len(STUB))
    # 2. retarget the call site (may straddle two records)
    write(CALL_SITE, NEW_CALL, expect=OLD_CALL)

    for payload, start, end in touched.values():
        reseal(payload, start, end)
    print(f"  {len(touched)} record(s) modified; message count and sizes unchanged")

    out = b"".join(rebuild_message(p, bytes(pl)) for p, pl in msgs)
    open(args.outfile, "wb").write(out)

    print(f"in  : {args.infile} ({len(data)} bytes, {len(msgs)} messages)")
    print(f"out : {args.outfile} ({len(out)} bytes, {len(msgs)} messages)")
    print(f"  0x{CALL_SITE:04X}: LCALL 0x{ROUTER:04X} -> LCALL 0x{STUB_ADDR:04X}")
    print(f"  0x{STUB_ADDR:04X}: {len(STUB)}-byte stub added ({STUB.hex(' ')})")
    return 0


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