#!/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.

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

SRC = sys.argv[1] if len(sys.argv) > 1 else "RIOv4_2.bin"
DST = sys.argv[2] if len(sys.argv) > 2 else "RIOv4_2_patched.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"

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}")