#!/usr/bin/env python3
"""Apply the RIO v4.2 reply-wedge fix to RIOv4_2.bin -> RIOv4_2_patched.bin.

Fix (see RIOv4_2-ANALYSIS.md): clear the reply-in-progress latch $2521 on
EVERY reply teardown, not just the $2522-gated success path.

  1. Give-up path: redirect $D9DD `JMP $DA2F` to a stub at free ROM $DFF0
     that clears $2521/$2522 then continues to $DA2F.
  2. Success path: make $DA00's clear of $2521/$2522 unconditional.

With --baud31250, additionally retune the SCI from 9600 to 31250 baud
(-> RIOv4_2_patched_31250.bin): the init at $D62A loads BAUD ($102B) with
$30 (prescale /13, divider /1 -> 2MHz E / (16*13) = 9615). $02 selects
prescale /1, divider /4 -> 2MHz / (16*4) = 31250 exactly. One operand
byte at $D62B. NOTE: 31250 is a non-standard rate — FTDI-class USB
adapters produce it exactly; classic 16550 UARTs cannot. The native
games still expect 9600, so a 31250 chip is bench/RIOJoy-only until the
games are patched.

Each edit asserts the exact original bytes first, so a wrong assumption
aborts instead of corrupting the image. Address == file offset.
"""
import sys, hashlib

# Baud retunes: flag -> (BAUD register value, ACK-wait ticks for ~12.8ms, name).
# The ACK-wait loop self-clocks in byte times (see edit 4), so each rate needs
# its own tick count to hold the same ~12.8ms wall-clock ACK grace window.
BAUD_OPTS = {
    "--baud31250":  (0x02, 0x28, "31250"),   # /1 /4  -> 31250  (40 ticks)
    "--baud62500":  (0x01, 0x50, "62500"),   # /1 /2  -> 62500  (80 ticks)
    "--baud125000": (0x00, 0xA0, "125000"),  # /1 /1  -> 125000 (160 ticks)
}
selected = [f for f in BAUD_OPTS if f in sys.argv]
assert len(selected) <= 1, f"pick one baud flag, got {selected}"
BAUD_VAL, ACKWAIT_VAL, BAUD_NAME = BAUD_OPTS[selected[0]] if selected else (None, None, None)

# --widen-ackwait: required semantics for 62500/125000 (implied there); optional
# for 31250 to keep the historical v1 image reproducible.
WIDEN_ACKWAIT = "--widen-ackwait" in sys.argv or BAUD_NAME in ("62500", "125000")
assert not ("--widen-ackwait" in sys.argv) or selected, \
    "--widen-ackwait only makes sense with a baud retune (the wait is byte-scaled)"

# --e0thresh[=N]: gate the E0 error display (see edit 5). Default N=5.
E0T = None
for a in sys.argv:
    if a.startswith("--e0thresh"):
        E0T = int(a.split("=", 1)[1]) if "=" in a else 5
assert E0T is None or 1 <= E0T <= 255, f"--e0thresh out of range: {E0T}"

# --checkrepaint: repaint the display after the CheckRequest self-test (edit 6).
CHECKREPAINT = "--checkrepaint" in sys.argv
assert not CHECKREPAINT or E0T is not None, \
    "--checkrepaint requires --e0thresh (it re-renders E0 via the gated $D5F2)"

# --reportversion=M.N (edit 7): the major.minor bytes the VersionReply returns.
REPORTVER = None
for a in sys.argv:
    if a.startswith("--reportversion="):
        REPORTVER = tuple(int(x) for x in a.split("=", 1)[1].split("."))
        assert len(REPORTVER) == 2 and all(0 <= v <= 127 for v in REPORTVER), \
            f"--reportversion must be M.N with 7-bit values: {REPORTVER}"

args = [a for a in sys.argv[1:] if not a.startswith("--")]
SRC = args[0] if len(args) > 0 else "RIOv4_2.bin"
if len(args) > 1:
    DST = args[1]
elif BAUD_NAME == "31250":
    DST = "RIOv4_2_patched_31250v2.bin" if WIDEN_ACKWAIT else "RIOv4_2_patched_31250.bin"
elif BAUD_NAME:
    DST = f"RIOv4_2_patched_{BAUD_NAME}.bin"
else:
    DST = "RIOv4_2_patched.bin"
if E0T is not None and len(args) <= 1:
    DST = DST[:-4] + f"_e0t{E0T}.bin"
if CHECKREPAINT and len(args) <= 1:
    DST = DST[:-4] + "_ckr.bin"

d = bytearray(open(SRC, "rb").read())
assert len(d) == 0x10000, f"expected 64KB image, got {len(d)}"
orig_sha = hashlib.sha256(d).hexdigest()
assert orig_sha == "60a88718835c654b6135dbec7721c40ef99dca07df2ad4b57eedeb24037a5f73", \
    f"unexpected source image {orig_sha}"

def patch(addr, expect, new):
    got = bytes(d[addr:addr+len(expect)])
    assert got == bytes(expect), (
        f"@${addr:04X}: expected {got.hex()} to be {bytes(expect).hex()}")
    assert len(new) == len(expect), "length mismatch"
    d[addr:addr+len(new)] = bytes(new)

# --- edit 1: give-up path redirect ---------------------------------------
# $D9DD  7E DA 2F   JMP $DA2F   ->   7E DF F0   JMP $DFF0
patch(0xD9DD, [0x7E, 0xDA, 0x2F], [0x7E, 0xDF, 0xF0])

# stub at $DFF0 (was erased $FF): CLR $2521; CLR $2522; JMP $DA2F
patch(0xDFF0, [0xFF]*8,
      [0x7F, 0x25, 0x21,   # CLR $2521
       0x7F, 0x25, 0x22,   # CLR $2522
       0x7E, 0xDA])        # JMP $DA2F (hi + first target byte)
patch(0xDFF8, [0xFF], [0x2F])   # JMP low byte

# --- edit 2: success teardown, unconditional clear -----------------------
# $DA21  B6 25 22  LDAA $2522
# $DA24  81 01     CMPA #$01
# $DA26  26 06     BNE  $DA2E
# $DA28  7F 25 21  CLR  $2521
# $DA2B  7F 25 22  CLR  $2522     (13 bytes $DA21-$DA2D; $DA2E RTS untouched)
# ->  CLR $2521 ; CLR $2522 ; NOP x7
patch(0xDA21,
      [0xB6,0x25,0x22, 0x81,0x01, 0x26,0x06, 0x7F,0x25,0x21, 0x7F,0x25,0x22],
      [0x7F,0x25,0x21, 0x7F,0x25,0x22, 0x01,0x01,0x01,0x01,0x01,0x01,0x01])
assert d[0xDA2E] == 0x39, "RTS at $DA2E must be intact"

# --- edit 3 (--baud31250 only): SCI 9600 -> 31250 --------------------------
# $D62A  86 30   LDAA #$30 (SCP=/13, SCR=/1)  ->  86 02 (SCP=/1, SCR=/4)
# then STAA $102B (BAUD). 2MHz E-clock: 2e6/(16*13)=9615 -> 2e6/(16*4)=31250.
if BAUD_VAL is not None:
    assert d[0xD62A] == 0x86 and bytes(d[0xD62C:0xD62F]) == bytes([0xB7,0x10,0x2B]), \
        "expected LDAA #imm ; STAA $102B at $D62A"
    patch(0xD62B, [0x30], [BAUD_VAL])

# --- edit 4 (--widen-ackwait; implied for 62500/125000): scale the ACK-wait --
# The no-ACK wait loop at $D9E0 self-clocks in BYTE TIMES: each tick transmits
# an IDLE ($FF) keep-alive and re-enters on that byte's TX-complete interrupt.
# The stock limit (4 ticks) = ~5.2ms at 9600 (USB ACK wins) but shrinks with
# baud (retry storm at 31250; bench 2026-07-18: framing 5655, Abandon 65).
# Each retune holds the window at ~12.8ms wall clock: 40/80/160 ticks.
#   $D9E6  81 04  CMPA #$04   ->   81 <ticks>
if WIDEN_ACKWAIT:
    assert d[0xD9E6] == 0x81 and bytes(d[0xD9E3:0xD9E6]) == bytes([0xB6,0x31,0x7B]), \
        "expected LDAA $317B ; CMPA #imm at $D9E3"
    patch(0xD9E7, [0x04], [ACKWAIT_VAL])

# --- edit 5 (--e0thresh=N): threshold-gate the E0 error display -----------
# Stock behavior: every increment of $3187 (TX-ring overflow), $3184 (reply
# retransmit) or $3185 (reply teardown/give-up) immediately calls $D5F2,
# which paints 'E0'+counters over the F0 banner — so ONE benign give-up
# shows E0000001 forever. Gate the render: skip unless any counter >= N.
# $D5F2 prologue (PSHX/PSHA) has already run when we take over its
# LDX #$2038; the cave may clobber A/X because both exits restore them —
# below-threshold leaves via the routine's own epilogue at $D61D
# (PULA/PULX/RTS), at-or-above resumes the render at $D5F7.
if E0T is not None:
    patch(0xD5F4, [0xCE, 0x20, 0x38], [0x7E, 0xE0, 0x00])  # LDX #$2038 -> JMP $E000
    cave = [0xB6, 0x31, 0x87,   # $E000  LDAA $3187  (ring overflows)
            0x81, E0T,          # $E003  CMPA #N
            0x24, 0x11,         # $E005  BCC  $E018  (>=N -> render)
            0xB6, 0x31, 0x84,   # $E007  LDAA $3184  (reply retransmits)
            0x81, E0T,          # $E00A  CMPA #N
            0x24, 0x0A,         # $E00C  BCC  $E018
            0xB6, 0x31, 0x85,   # $E00E  LDAA $3185  (teardowns/give-ups)
            0x81, E0T,          # $E011  CMPA #N
            0x24, 0x03,         # $E013  BCC  $E018
            0x7E, 0xD6, 0x1D,   # $E015  JMP  $D61D  (below threshold: epilogue)
            0xCE, 0x20, 0x38,   # $E018  LDX  #$2038 (displaced instruction)
            0x7E, 0xD5, 0xF7]   # $E01B  JMP  $D5F7  (resume render)
    patch(0xE000, [0xFF] * len(cave), cave)

# --- edit 6 (--checkrepaint): repaint the display after the check self-test --
# The CheckRequest handler ($C5A6) renders its five-channel encoder sweep to
# the cockpit display (it sets the test-display flag $2421 for the duration)
# and returns without repainting, leaving a stale '04 00 <ch4 count>' frame
# (idle: 04000000) after every status check. Hook its final TestModeChange
# notify ($C5E4 JSR $C70E) through a cave that also repaints F0000000
# ($CB7C) and then re-renders the E0 error readout via $D5F2 — which, with
# the edit-5 gate, paints only when a counter is at/over threshold. Display
# after any check: F0000000 when healthy, E0... when not.
if CHECKREPAINT:
    patch(0xC5E4, [0xBD, 0xC7, 0x0E], [0xBD, 0xE0, 0x20])  # JSR $C70E -> JSR $E020
    cave6 = [0xBD, 0xC7, 0x0E,   # $E020  JSR $C70E  (original notify, kept)
             0xBD, 0xCB, 0x7C,   # $E023  JSR $CB7C  (repaint F0000000)
             0xBD, 0xD5, 0xF2,   # $E026  JSR $D5F2  (E0 readout iff >= threshold)
             0x39]               # $E029  RTS
    patch(0xE020, [0xFF] * len(cave6), cave6)

# --- edit 7 (--reportversion=M.N): version bytes in the VersionReply --------
# The reply builder at $C6EA loads the payload with LDAA #$04 ($C6F7) and
# LDAA #$02 ($C6FE) — hardcoded "4.2". Nothing host-side validates the value
# (legacy app prints it, RIOJoy parses it, native games ignore it — verified
# by Cyd 2026-07-19), so the christened firmware can announce itself.
if REPORTVER is not None:
    assert d[0xC6F7] == 0x86 and d[0xC6FE] == 0x86, "expected LDAA #imm pair at $C6F7/$C6FE"
    patch(0xC6F8, [0x04], [REPORTVER[0]])
    patch(0xC6FF, [0x02], [REPORTVER[1]])

open(DST, "wb").write(d)
new_sha = hashlib.sha256(d).hexdigest()
# byte-diff report
diffs = [(a, orig, d[a]) for a, orig in
         enumerate(open(SRC,"rb").read()) if d[a] != orig]
print(f"source : {SRC}  sha256 {orig_sha}")
print(f"patched: {DST}  sha256 {new_sha}")
print(f"{len(diffs)} bytes changed:")
for a, o, n in diffs:
    print(f"  ${a:04X}: {o:02X} -> {n:02X}")