#!/usr/bin/env python3
# patch_c2.py -- apply the USB-1->CV patch to the QuNexus via the C2 interface,
# surgically: erase + reprogram only the TWO flash pages the patch touches.
#
# The patch (see ../patch_usb1_to_cv.py, STUB_ADDR=0x8126):
#   * page 0x8000-0x81FF: a 23-byte stub written into 0xFF linker padding at
#     0x8126 (verified 0xFF in stock). Routes cable-0 events to the CV ring
#     after the normal router call.
#   * page 0xDE00-0xDFFF: retarget the LCALL at 0xDF28 from 0xA57B to 0x8126.
#
# SAFETY / ORDERING: the stub page is written and verified FIRST, the retarget
# page SECOND. So at no intermediate point does the retarget reference a stub
# that isn't there. If the stub-page step fails, we STOP and the device is left
# stock (working, unpatched). If the retarget-page step fails, the stub is in
# place but unreferenced -> also stock behaviour. The device can never be bricked
# mid-process. The bootloader (0x0000-0x23FF) and lock/reserved (0xFA00+) are
# never touched; only app-region pages 0x40 and 0x6F are erased/written.
import os, sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, ".."))     # for patch_usb1_to_cv
import cdc
from patch_usb1_to_cv import STUB, STUB_ADDR, CALL_SITE, NEW_CALL, OLD_CALL, ROUTER

STOCK_BIN = os.path.join(HERE, "..", "out", "QuNexus_Firmware_v2.2.1.bin")
STOCK_BASE = 0x2400
PAGE = 512

STUB_PAGE  = 0x8000          # page 0x40
CALL_PAGE  = 0xDE00          # page 0x6F


def stock_page(addr):
    s = open(STOCK_BIN, "rb").read()
    o = addr - STOCK_BASE
    return bytearray(s[o:o + PAGE])


def patch_page(page_addr, edits):
    """edits: list of (addr, bytes). Returns patched 512-byte page."""
    p = stock_page(page_addr)
    for addr, data in edits:
        off = addr - page_addr
        assert 0 <= off and off + len(data) <= PAGE, f"edit 0x{addr:04x} outside page 0x{page_addr:04x}"
        p[off:off + len(data)] = data
    return p


def block_read(fd, addr, n=256):
    lc = n if n < 256 else 0
    r = cdc.cmd(fd, f"read {addr:04x} {lc}", 12).strip()
    if r.startswith("data "):
        try:
            b = bytes.fromhex(r.split()[3])
            if len(b) == n:
                return b
        except (IndexError, ValueError):
            pass
    return None


def read_page(fd, page_addr):
    lo = block_read(fd, page_addr, 256)
    hi = block_read(fd, page_addr + 256, 256)
    if lo is None or hi is None:
        return None
    return lo + hi


def erase_page(fd, page_addr):
    return cdc.cmd(fd, f"pe {page_addr // PAGE:x}", 25).strip()


def write_block(fd, addr, data):
    lc = len(data) if len(data) < 256 else 0
    return cdc.cmd(fd, f"bw {addr:04x} {lc} " + data.hex(), 25).strip()


def write_page(fd, page_addr, data):
    r1 = write_block(fd, page_addr, data[:256])
    r2 = write_block(fd, page_addr + 256, data[256:])
    return r1, r2


def fail(msg):
    print(f"\n[FAIL] {msg}")
    print("       Device left in a WORKING state (retarget not written, or stub")
    print("       written but unreferenced). Re-run to retry.")
    return 1


def do_page(fd, label, page_addr, patched, expect_stock_first=True):
    """Full verified erase+write of one patched page. Returns True on success."""
    pgnum = page_addr // PAGE
    print(f"\n[{label}] page 0x{page_addr:04x}-0x{page_addr+PAGE-1:04x} (page 0x{pgnum:02x})")

    # 1. read current page; confirm it's currently stock (clean baseline)
    cur = read_page(fd, page_addr)
    if cur is None:
        print("  err: could not read current page"); return False
    stock = stock_page(page_addr)
    if cur != stock:
        nd = sum(1 for i in range(PAGE) if cur[i] != stock[i])
        print(f"  WARNING: current page is NOT stock ({nd} bytes differ).")
        if expect_stock_first:
            print("  Refusing to patch a non-stock page (unexpected state).")
            return False
    print("  current page == stock  OK")

    # 2. erase
    r = erase_page(fd, page_addr)
    print(f"  pe: {r}")
    if r != "ok":
        print("  err: page erase failed"); return False

    # 3. verify erased
    er = read_page(fd, page_addr)
    if er is None or sum(1 for b in er if b == 0xFF) != PAGE:
        print("  err: page did not erase to all 0xFF"); return False
    print("  erased: all 0xFF  OK")

    # 4. write patched bytes
    r1, r2 = write_page(fd, page_addr, patched)
    print(f"  bw lo: {r1}")
    print(f"  bw hi: {r2}")
    if not (r1.startswith("ok bw") and r2.startswith("ok bw")):
        print("  err: block write failed"); return False

    # 5. verify == patched
    got = read_page(fd, page_addr)
    if got is None:
        print("  err: could not read back page"); return False
    if got != bytes(patched):
        diffs = [i for i in range(PAGE) if got[i] != patched[i]]
        print(f"  err: write-back MISMATCH ({len(diffs)} bytes; first {diffs[:8]})")
        return False
    # confirm only the intended bytes changed vs stock
    intended = [i for i in range(PAGE) if patched[i] != stock[i]]
    actual   = [i for i in range(PAGE) if got[i] != stock[i]]
    if intended != actual:
        print(f"  err: unintended bytes changed. intended {len(intended)}, actual {len(actual)}")
        return False
    print(f"  verified == patched; {len(intended)} byte(s) changed vs stock  OK")
    return True


def main():
    # build the two patched pages from stock + patch constants
    stub_page = patch_page(STUB_PAGE, [(STUB_ADDR, STUB)])
    call_page = patch_page(CALL_PAGE, [(CALL_SITE, NEW_CALL)])

    # sanity: confirm the stub lands on 0xFF and the call site is the old LCALL
    s = stock_page(STUB_PAGE)
    assert all(s[STUB_ADDR - STUB_PAGE + k] == 0xFF for k in range(len(STUB))), "stub site not 0xFF in stock"
    c = stock_page(CALL_PAGE)
    assert bytes(c[CALL_SITE - CALL_PAGE:CALL_SITE - CALL_PAGE + 3]) == OLD_CALL, "call site not old LCALL in stock"

    print(f"[patch] stub  @0x{STUB_ADDR:04x} ({len(STUB)} B): {STUB.hex(' ')}")
    print(f"[patch] call @0x{CALL_SITE:04x}: {OLD_CALL.hex(' ')} -> {NEW_CALL.hex(' ')}  (LCALL 0x{ROUTER:04X} -> LCALL 0x{STUB_ADDR:04X})")
    print(f"[patch] pages to modify: 0x{STUB_PAGE:04x} (stub, written FIRST) and 0x{CALL_PAGE:04x} (retarget, written SECOND)")

    fd = cdc.open_port()
    cdc.drain(fd, 0.4)

    def c(s, t=15): return cdc.cmd(fd, s, t).strip()

    print("\n[init] piinit + wsetup")
    r = c("piinit", 12); print("  piinit:", r)
    if r != "ok piinit": return fail("piinit failed")
    r = c("wsetup", 12); print("  wsetup:", r)
    if r != "ok": return fail("wsetup failed")

    # --- stub page FIRST ---
    if not do_page(fd, "STUB", STUB_PAGE, stub_page):
        return fail("stub page step failed -- retarget NOT written; device is stock/working")

    # --- retarget page SECOND ---
    if not do_page(fd, "CALL", CALL_PAGE, call_page):
        return fail("retarget page step failed -- stub is written but unreferenced; device is stock/working")

    # --- reset so the patched app boots ---
    print("\n[reset] booting patched app")
    c("reset", 5)
    print("  ok reset")

    print("\n[DONE] patch applied & verified. 0xDF28 now LCALLs 0x8126, which runs the")
    print("       normal router then re-routes cable-0 (USB-1) events to the USB-3/CV ring.")
    print("       Plug the QuNexus into the Mac and test: send MIDI to USB port 1 and")
    print("       confirm CV output responds (in addition to the normal Control Surface path).")
    return 0


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