#!/usr/bin/env python3
# reflash_page.py -- surgical recovery: erase + reprogram ONE flash page to
# undo the bad patch, restoring stock 2.2.1 code at the LCALL site 0xDF28.
#
# What it does: page 0xDE00-0xDFFF (512 B) contains 0xDF28 (the LCALL the patch
# retargeted to empty 0xE8F2). We erase that one page and reprogram it with the
# stock bytes, restoring 0xDF29 = A5 7B (LCALL 0xA57B). This un-breaks the
# USB-MIDI dispatcher, so SysEx / bootloader entry work again.
#
# SAFETY: only page 0x6F (0xDE00) is erased, and only 0xDE00/0xDF00 written.
# The firmware hard-guards all erase/write to the app region 0x2400-0xF9FF; the
# bootloader (0x0000-0x23FF) and lock/reserved (0xFA00+) are unreachable.
# Every step is verified; the script aborts on any mismatch.
import os, sys, time, select, termios

DEV = os.environ.get("C2PROBE_DEV", "/dev/cu.usbmodem2101")
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import cdc

STOCK_BIN = os.path.join(HERE, "..", "out", "QuNexus_Firmware_v2.2.1.bin")
STOCK_BASE = 0x2400
PAGE_ADDR = 0xDE00
PAGE_LEN = 512
PAGE_NUM = PAGE_ADDR // 512        # 0x6F
LCALL_ADDR = 0xDF28

def main():
    stock = open(STOCK_BIN, "rb").read()
    page = stock[PAGE_ADDR - STOCK_BASE : PAGE_ADDR - STOCK_BASE + PAGE_LEN]
    assert len(page) == PAGE_LEN, "stock page incomplete"
    expect_call = stock[LCALL_ADDR - STOCK_BASE : LCALL_ADDR - STOCK_BASE + 3]
    print(f"[stock] page 0x{PAGE_ADDR:04x}-0x{PAGE_ADDR+PAGE_LEN-1:04x} ready")
    print(f"[stock] 0x{LCALL_ADDR:04x} = {expect_call.hex(' ')} (target: restore this)")

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

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

    # 1. halt core + flash-programming SFR setup
    print("\n[1] 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")

    # 2. read current page, confirm the patch is present (0xDF29 = e8 f2)
    print("\n[2] read current page (confirm patch present)")
    cur_lo = block_read(fd, 0xDE00, 256)
    cur_hi = block_read(fd, 0xDF00, 256)
    if cur_lo is None or cur_hi is None: return fail("could not read current page")
    cur = cur_lo + cur_hi
    print(f"   0x{LCALL_ADDR:04x} now = {cur[LCALL_ADDR-PAGE_ADDR:LCALL_ADDR-PAGE_ADDR+3].hex(' ')}")
    if cur[LCALL_ADDR-PAGE_ADDR+1:LCALL_ADDR-PAGE_ADDR+3] != b"\xe8\xf2":
        print("   WARNING: 0xDF29 is not e8 f2 -- patch may already be undone")

    # 3. erase the page
    print(f"\n[3] page erase page 0x{PAGE_NUM:02x} (0x{PAGE_ADDR:04x})")
    r = c(f"pe {PAGE_NUM:x}", 20)
    print("   pe:", r)
    if r != "ok": return fail(f"page erase failed: {r}")

    # 4. verify the page is now all 0xFF
    print("\n[4] verify page erased (all 0xFF)")
    er_lo = block_read(fd, 0xDE00, 256)
    er_hi = block_read(fd, 0xDF00, 256)
    if er_lo is None or er_hi is None: return fail("could not read erased page")
    erased = er_lo + er_hi
    nff = sum(1 for b in erased if b == 0xFF)
    print(f"   0xFF count: {nff}/512")
    if nff != 512:
        print("   erased page:", erased.hex())
        return fail("page erase did NOT yield all 0xFF -- aborting before write")

    # 5. program the stock bytes (two 256-byte block writes)
    print("\n[5] block write stock bytes")
    r = c("bw de00 0 " + page[:256].hex(), 20);  print("   bw de00:", r)
    if not r.startswith("ok bw"): return fail(f"bw de00 failed: {r}")
    r = c("bw df00 0 " + page[256:].hex(), 20);  print("   bw df00:", r)
    if not r.startswith("ok bw"): return fail(f"bw df00 failed: {r}")

    # 6. verify the page now matches stock
    print("\n[6] verify page == stock")
    v_lo = block_read(fd, 0xDE00, 256)
    v_hi = block_read(fd, 0xDF00, 256)
    if v_lo is None or v_hi is None: return fail("could not read back page")
    got = v_lo + v_hi
    if got == page:
        print(f"   MATCH -- 0x{LCALL_ADDR:04x} = {got[LCALL_ADDR-PAGE_ADDR:LCALL_ADDR-PAGE_ADDR+3].hex(' ')}")
    else:
        diffs = [i for i in range(512) if got[i] != page[i]]
        print(f"   MISMATCH: {len(diffs)} bytes differ; first: {diffs[:8]}")
        return fail("write-back did not match stock")

    # 7. reset the device so it boots the restored app
    print("\n[7] reset device (boot restored app)")
    c("reset", 5)
    print("   ok reset")

    print("\n[DONE] page 0xDE00 restored to stock. The LCALL at 0xDF28 now targets")
    print("       0xA57B again, so the USB-MIDI dispatcher is intact. The device")
    print("       should enumerate and respond to MIDI / SysEx. Test it, then if")
    print("       you want a full factory-fresh image, reflash v2.2.1 over USB.")
    return 0

def block_read(fd, addr, n):
    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 fail(msg):
    print(f"\n[FAIL] {msg}")
    print("       The device is NO worse than before (bootloader untouched).")
    print("       Re-run this script to retry.")
    return 1

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