#!/usr/bin/env python3
"""
patch_rfs_mz5z008.py - Patch MZ-5Z008 binary to add RFS (ROM Filing System) device support.

Adds a $RFS device to the device chain, enabling:
  DIR "RFS:"   - List SD card directory (drive 0)
  DIR "RFS2:"  - List SD card directory (drive 1)
  LOAD "RFS:"  - Load from SD card
  SAVE "RFS:"  - Save to SD card

The RFS device code is appended at 85FFH (end of binary).
The $QD device's link pointer is patched to chain: $QD -> $RFS -> $USR.
TEXTST and MEMCLI are moved past the RFS code.

No DIR hook is needed — MZ-5Z008 uses SVC .DIR which calls FDIR in the monitor.
FDIR calls ZRO in a loop to enumerate files for SEQ devices.

Requires the picoZ80 RFS ROM at E800H (virtual hardware in the Z80 emulator).

Usage:
  python3 patch_rfs_mz5z008.py [input_binary] [output_binary]

  input_binary:  Original MZ-5Z008 binary (default: mz5z008.bin)
  output_binary: Patched binary with RFS support (default: mz5z008_rfs.bin)
"""

import struct
import sys
import os

# --- Constants ---
BINARY_SIZE     = 34303              # Expected MZ-5Z008 binary size (0x85FF bytes)
RFS_BASE        = 0x85FF             # Append RFS code at end of binary
QD_LINK         = 0x38A2             # $QD device descriptor (link pointer at +0)
USR_ADDR        = 0x0D5E             # $USR device descriptor address
TEXTST_INIT     = 0x4000             # "LD HL,85FFH" in BASIC init
MEMCLI_INIT     = 0x4020             # "RST 18H / DB 20H / EX DE,HL" in MEMCLI start
VERSION_OFFSET  = 0x8536             # Version letter in "V1.0A" string (ASCII)
VERSION_LETTER  = ord('B')           # Change to 'B' for patched builds

# MZ-5Z008 addresses
ELMD            = 0x1000
ELMD1           = 0x1001
IBUFE           = 0x10F0
ZCH             = 0x1046
SDDRIVENO       = 0x1028
CMTFILENO       = 0x1029
BREAKZ          = 0x00E5
HDRLEN          = 24

# RFS ROM API addresses
UROMADDR        = 0xE800
RFSJMPTABLE     = UROMADDR + 0xB0
RFSF_RDINF      = RFSJMPTABLE + 0x06   # E8B6
RFSF_RDDATA     = RFSJMPTABLE + 0x09   # E8B9
RFSF_WRINF      = RFSJMPTABLE + 0x0C   # E8BC
RFSF_WRDATA     = RFSJMPTABLE + 0x0F   # E8BF


def le16(val):
    """Encode 16-bit value as little-endian bytes."""
    return bytes([val & 0xFF, (val >> 8) & 0xFF])


def build_rfs_code():
    """
    Build the RFS device module as a byte array.
    Returns (code_bytes, symbol_table) where symbol_table maps label names to addresses.

    Memory banking: MZ-5Z008 BASIC runs with upper area (D000-FFFF) as DRAM
    (via OUT E1H). The RFS ROM at E800H is hidden. Each ROM call must:
      1. Save SP (BASIC stack is likely in D000-FFFF range)
      2. Switch SP to temp stack in safe area (1000-CFFF)
      3. DI (prevent interrupts while banked)
      4. XOR A — clear A for bank select
      5. OUT (E3H),A — switch upper area to MMIO (exposes RFS ROM)
      6. CALL TRAIN_UPCNTR — read BNKCTRLRST 15 times to train the register
         latch (emulation doesn't auto-train from instruction fetches like
         real hardware does; upCntr must reach 15 before register writes work)
      7. LD (BNKSELUSER),A — select User ROM bank 0 (entry points at E8B0H+)
      8. CALL ROM function
      9. Save flags (carry = error indicator)
      10. OUT (E1H),A — switch upper area back to DRAM
      11. EI
      12. Restore flags and SP
    """
    code = bytearray()
    symbols = {}
    base = RFS_BASE

    def addr():
        return base + len(code)

    def emit(*args):
        for b in args:
            if isinstance(b, (bytes, bytearray)):
                code.extend(b)
            elif isinstance(b, int):
                code.append(b & 0xFF)
            else:
                raise ValueError(f"Unknown type: {type(b)}")

    def emit16(val):
        emit(le16(val))

    # Port constants for memory banking
    PORT_DRAM_UPPER = 0xE1    # OUT (E1H),A = upper area -> DRAM
    PORT_MMIO_UPPER = 0xE3    # OUT (E3H),A = upper area -> Memory Mapped I/O
    BNKSELUSER      = 0xEFFE  # Memory-mapped register: select User ROM bank (LD (0EFFEh),A)
    BNKCTRLRST      = 0xEFF8  # Bank control reset register (safe to read for upCntr training)

    # =====================================================================
    # $RFS Device Descriptor (29 bytes)
    # =====================================================================
    symbols['$RFS'] = addr()

    # Reserve 29 bytes, fill after we know handler addresses
    desc_offset = len(code)
    code.extend(b'\x00' * 29)

    # =====================================================================
    # RFSFREE - Free space query (not supported, return carry)
    # =====================================================================
    symbols['RFSFREE'] = addr()
    emit(0xAF)                     # XOR A
    emit(0x37)                     # SCF
    emit(0xC9)                     # RET

    # =====================================================================
    # RFSINIT - Initialize RFS device
    # =====================================================================
    symbols['RFSINIT'] = addr()
    emit(0x3A, *le16(ZCH))        # LD A,(ZCH)
    emit(0xC6, 0x30)              # ADD A,30H        ; '0'
    emit(0x32, *le16(SDDRIVENO))  # LD (SDDRIVENO),A
    emit(0xAF)                     # XOR A
    emit(0x32, *le16(CMTFILENO))  # LD (CMTFILENO),A
    # Clear first-call flag
    rfsinit_first_patch = len(code)
    emit(0x32, 0x00, 0x00)        # LD (RFS_FIRST),A  [patched below]
    emit(0xC9)                     # RET

    # =====================================================================
    # RFSSTRT - Start/Open device (ZCL/ZSTRT with A=0 means start)
    # =====================================================================
    symbols['RFSSTRT'] = addr()
    emit(0xB7)                     # OR A
    emit(0xC0)                     # RET NZ            ; A!=0 = close, just return
    emit(0x3A, *le16(ZCH))        # LD A,(ZCH)
    emit(0xC6, 0x30)              # ADD A,30H
    emit(0x32, *le16(SDDRIVENO))  # LD (SDDRIVENO),A
    emit(0xAF)                     # XOR A
    emit(0x32, *le16(CMTFILENO))  # LD (CMTFILENO),A
    emit(0xC9)                     # RET

    # =====================================================================
    # CALL_ROM - Helper to call RFS ROM with memory banking
    #
    # The RFS ROM is at E800H, hidden behind DRAM when BASIC runs.
    # This helper banks in MMIO, calls the ROM, banks back to DRAM.
    # BASIC's stack is likely in D000-FFFF, so we use a temp stack
    # in the safe 1000-CFFF region.
    #
    # Entry: Call address embedded as CALL instruction (4 wrappers below)
    #        DE, other regs set up for the ROM function
    # Exit:  Carry flag preserved from ROM call
    #        All other regs preserved
    # =====================================================================
    # We implement 4 specific wrappers instead of a generic helper,
    # to avoid needing indirect call machinery.

    # --- CALL_RDINF wrapper ---
    symbols['CALL_RDINF'] = addr()
    call_rdinf_save_sp = len(code)
    emit(0xED, 0x73, 0x00, 0x00)  # LD (SAVE_SP),SP   [patched below]
    call_rdinf_set_sp = len(code)
    emit(0x31, 0x00, 0x00)        # LD SP,TEMP_STACK_TOP [patched below]
    emit(0xF3)                     # DI
    emit(0xAF)                     # XOR A              ; A = 0
    emit(0xD3, PORT_MMIO_UPPER)   # OUT (E3H),A        ; upper -> MMIO
    call_rdinf_train = len(code)
    emit(0xCD, 0x00, 0x00)        # CALL TRAIN_UPCNTR  ; train register latch [patched below]
    emit(0x32, *le16(BNKSELUSER)) # LD (BNKSELUSER),A  ; select User ROM bank 0
    emit(0xCD, *le16(RFSF_RDINF)) # CALL RFSF_RDINF
    emit(0xF5)                     # PUSH AF            ; save flags (on temp stack)
    emit(0xD3, PORT_DRAM_UPPER)   # OUT (E1H),A        ; upper -> DRAM
    emit(0xFB)                     # EI
    emit(0xF1)                     # POP AF             ; restore flags
    call_rdinf_rest_sp = len(code)
    emit(0xED, 0x7B, 0x00, 0x00)  # LD SP,(SAVE_SP)   [patched below]
    emit(0xC9)                     # RET

    # --- CALL_RDDATA wrapper ---
    symbols['CALL_RDDATA'] = addr()
    call_rddata_save_sp = len(code)
    emit(0xED, 0x73, 0x00, 0x00)  # LD (SAVE_SP),SP
    call_rddata_set_sp = len(code)
    emit(0x31, 0x00, 0x00)        # LD SP,TEMP_STACK_TOP
    emit(0xF3)                     # DI
    emit(0xAF)                     # XOR A              ; A = 0
    emit(0xD3, PORT_MMIO_UPPER)   # OUT (E3H),A
    call_rddata_train = len(code)
    emit(0xCD, 0x00, 0x00)        # CALL TRAIN_UPCNTR  [patched below]
    emit(0x32, *le16(BNKSELUSER)) # LD (BNKSELUSER),A  ; select User ROM bank 0
    emit(0xCD, *le16(RFSF_RDDATA))# CALL RFSF_RDDATA
    emit(0xF5)                     # PUSH AF
    emit(0xD3, PORT_DRAM_UPPER)   # OUT (E1H),A
    emit(0xFB)                     # EI
    emit(0xF1)                     # POP AF
    call_rddata_rest_sp = len(code)
    emit(0xED, 0x7B, 0x00, 0x00)  # LD SP,(SAVE_SP)
    emit(0xC9)                     # RET

    # --- CALL_WRINF wrapper ---
    symbols['CALL_WRINF'] = addr()
    call_wrinf_save_sp = len(code)
    emit(0xED, 0x73, 0x00, 0x00)  # LD (SAVE_SP),SP
    call_wrinf_set_sp = len(code)
    emit(0x31, 0x00, 0x00)        # LD SP,TEMP_STACK_TOP
    emit(0xF3)                     # DI
    emit(0xAF)                     # XOR A              ; A = 0
    emit(0xD3, PORT_MMIO_UPPER)   # OUT (E3H),A
    call_wrinf_train = len(code)
    emit(0xCD, 0x00, 0x00)        # CALL TRAIN_UPCNTR  [patched below]
    emit(0x32, *le16(BNKSELUSER)) # LD (BNKSELUSER),A  ; select User ROM bank 0
    emit(0xCD, *le16(RFSF_WRINF)) # CALL RFSF_WRINF
    emit(0xF5)                     # PUSH AF
    emit(0xD3, PORT_DRAM_UPPER)   # OUT (E1H),A
    emit(0xFB)                     # EI
    emit(0xF1)                     # POP AF
    call_wrinf_rest_sp = len(code)
    emit(0xED, 0x7B, 0x00, 0x00)  # LD SP,(SAVE_SP)
    emit(0xC9)                     # RET

    # --- CALL_WRDATA wrapper ---
    symbols['CALL_WRDATA'] = addr()
    call_wrdata_save_sp = len(code)
    emit(0xED, 0x73, 0x00, 0x00)  # LD (SAVE_SP),SP
    call_wrdata_set_sp = len(code)
    emit(0x31, 0x00, 0x00)        # LD SP,TEMP_STACK_TOP
    emit(0xF3)                     # DI
    emit(0xAF)                     # XOR A              ; A = 0
    emit(0xD3, PORT_MMIO_UPPER)   # OUT (E3H),A
    call_wrdata_train = len(code)
    emit(0xCD, 0x00, 0x00)        # CALL TRAIN_UPCNTR  [patched below]
    emit(0x32, *le16(BNKSELUSER)) # LD (BNKSELUSER),A  ; select User ROM bank 0
    emit(0xCD, *le16(RFSF_WRDATA))# CALL RFSF_WRDATA
    emit(0xF5)                     # PUSH AF
    emit(0xD3, PORT_DRAM_UPPER)   # OUT (E1H),A
    emit(0xFB)                     # EI
    emit(0xF1)                     # POP AF
    call_wrdata_rest_sp = len(code)
    emit(0xED, 0x7B, 0x00, 0x00)  # LD SP,(SAVE_SP)
    emit(0xC9)                     # RET

    # =====================================================================
    # RFSRO - Read Open (read next file header from SD for DIR enumeration)
    #
    # Called by FDIR in a loop. Each call reads the next file header.
    # On first call after RFSINIT, resets the file counter.
    # Returns: NC = success (caller's buffer filled), CY = no more files
    #
    # The sequential directory scanner (L1BFD) passes HL = 11A4H (KEYBUF)
    # and expects ZRO to write the directory entry there.  After ZRO
    # returns, L1C39 copies 11A4H → ELMD (1000H).  We must NOT write to
    # ELMD here — that would clobber the search filename at 1001H.
    # =====================================================================
    KEYBUF = 0x11A4

    symbols['RFSRO'] = addr()
    emit(0xE5)                     # PUSH HL
    emit(0xD5)                     # PUSH DE
    emit(0xC5)                     # PUSH BC
    # Check first-call flag
    rfsro_first_load = len(code)
    emit(0x3A, 0x00, 0x00)        # LD A,(RFS_FIRST)  [patched below]
    emit(0xB7)                     # OR A
    rfsro_jrnz = len(code)
    emit(0x20, 0x00)              # JR NZ,RFSRO1      [patched below]
    # First call: setup drive, reset file counter
    emit(0x3A, *le16(ZCH))        # LD A,(ZCH)
    emit(0xC6, 0x30)              # ADD A,30H
    emit(0x32, *le16(SDDRIVENO))  # LD (SDDRIVENO),A
    emit(0xAF)                     # XOR A
    emit(0x32, *le16(CMTFILENO))  # LD (CMTFILENO),A
    emit(0x3C)                     # INC A             ; A = 1
    rfsro_first_store = len(code)
    emit(0x32, 0x00, 0x00)        # LD (RFS_FIRST),A  [patched below]

    symbols['RFSRO1'] = addr()
    # Patch the JR NZ offset
    code[rfsro_jrnz + 1] = (addr() - (base + rfsro_jrnz + 2)) & 0xFF

    # Call RFS ROM (with banking) to get next file
    rfscr_load = len(code)
    emit(0x11, 0x00, 0x00)        # LD DE,RFSCR       [patched below]
    emit(0xCD, *le16(0x0000))     # CALL CALL_RDINF   [patched below]
    rfsro_call_rdinf = len(code) - 2  # offset of address in CALL
    rfsro_jrc = len(code)
    emit(0x38, 0x00)              # JR C,RFSROERR     [patched below]
    # Copy IBUFE -> KEYBUF (L1BFD's read buffer at 11A4H)
    emit(0x21, *le16(IBUFE))      # LD HL,IBUFE
    emit(0x11, *le16(KEYBUF))     # LD DE,11A4H       ; scanner's buffer
    emit(0x01, *le16(HDRLEN))     # LD BC,HDRLEN
    emit(0xED, 0xB0)              # LDIR
    emit(0xC1)                     # POP BC
    emit(0xD1)                     # POP DE
    emit(0xE1)                     # POP HL
    emit(0xAF)                     # XOR A             ; NC = success
    emit(0xC9)                     # RET

    symbols['RFSROERR'] = addr()
    code[rfsro_jrc + 1] = (addr() - (base + rfsro_jrc + 2)) & 0xFF
    emit(0xC1)                     # POP BC
    emit(0xD1)                     # POP DE
    emit(0xE1)                     # POP HL
    emit(0x37)                     # SCF               ; CY = no more files
    emit(0xC9)                     # RET

    # =====================================================================
    # RFSWO - Write Open (write file header to SD)
    # =====================================================================
    symbols['RFSWO'] = addr()
    emit(0xE5)                     # PUSH HL
    emit(0xD5)                     # PUSH DE
    emit(0xC5)                     # PUSH BC
    # Copy ELMD -> IBUFE
    emit(0x21, *le16(ELMD))       # LD HL,ELMD
    emit(0x11, *le16(IBUFE))      # LD DE,IBUFE
    emit(0x01, *le16(HDRLEN))     # LD BC,HDRLEN
    emit(0xED, 0xB0)              # LDIR
    # Set drive number
    emit(0x3A, *le16(ZCH))        # LD A,(ZCH)
    emit(0xC6, 0x30)              # ADD A,30H
    emit(0x32, *le16(SDDRIVENO))  # LD (SDDRIVENO),A
    # Call RFS ROM (with banking)
    emit(0x11, *le16(IBUFE + 1))  # LD DE,IBUFE+1 (NAME field)
    emit(0xCD, *le16(0x0000))     # CALL CALL_WRINF   [patched below]
    rfswo_call_wrinf = len(code) - 2
    emit(0xC1)                     # POP BC
    emit(0xD1)                     # POP DE
    emit(0xE1)                     # POP HL
    emit(0xC9)                     # RET

    # =====================================================================
    # RFSINP - Input (read file data from SD)
    # =====================================================================
    symbols['RFSINP'] = addr()
    emit(0xE5)                     # PUSH HL
    emit(0xD5)                     # PUSH DE
    emit(0xC5)                     # PUSH BC
    # Copy ELMD -> IBUFE (ROM needs DTADR/SIZE)
    emit(0x21, *le16(ELMD))       # LD HL,ELMD
    emit(0x11, *le16(IBUFE))      # LD DE,IBUFE
    emit(0x01, *le16(HDRLEN))     # LD BC,HDRLEN
    emit(0xED, 0xB0)              # LDIR
    emit(0xCD, *le16(0x0000))     # CALL CALL_RDDATA  [patched below]
    rfsinp_call_rddata = len(code) - 2
    emit(0xC1)                     # POP BC
    emit(0xD1)                     # POP DE
    emit(0xE1)                     # POP HL
    emit(0xC9)                     # RET

    # =====================================================================
    # RFSOUT - Output (write file data to SD)
    # =====================================================================
    symbols['RFSOUT'] = addr()
    emit(0xE5)                     # PUSH HL
    emit(0xD5)                     # PUSH DE
    emit(0xC5)                     # PUSH BC
    # Copy ELMD -> IBUFE
    emit(0x21, *le16(ELMD))       # LD HL,ELMD
    emit(0x11, *le16(IBUFE))      # LD DE,IBUFE
    emit(0x01, *le16(HDRLEN))     # LD BC,HDRLEN
    emit(0xED, 0xB0)              # LDIR
    emit(0xCD, *le16(0x0000))     # CALL CALL_WRDATA  [patched below]
    rfsout_call_wrdata = len(code) - 2
    emit(0xC1)                     # POP BC
    emit(0xD1)                     # POP DE
    emit(0xE1)                     # POP HL
    emit(0xC9)                     # RET

    # =====================================================================
    # TRAIN_UPCNTR - Train the RFS register latch (upCntr) to 15
    #
    # On real hardware, any bus cycle to E800-EFFF increments the upCntr
    # latch via the hardware address decoder. After 15 accesses the
    # control registers (EFF8-EFFF) become writable.
    #
    # In picoZ80 emulation, only explicit reads/writes to the register
    # addresses (EFF8-EFFF) increment upCntr. Normal opcode fetches
    # from E800+ do NOT train the counter. This subroutine explicitly
    # reads BNKCTRLRST 15 times to force-train the latch.
    #
    # Entry: MMIO must be mapped (OUT E3H done)
    # Exit:  A = 0 (ready for LD (BNKSELUSER),A)
    #        HL, B modified (scratch at this point in wrappers)
    # =====================================================================
    symbols['TRAIN_UPCNTR'] = addr()
    emit(0x21, *le16(BNKCTRLRST)) # LD HL,EFF8H     ; BNKCTRLRST address
    emit(0x06, 15)                 # LD B,15         ; 15 reads to train
    train_loop = len(code)
    emit(0x7E)                     # LD A,(HL)       ; read triggers handler, upCntr++
    emit(0x10)                     # DJNZ
    emit((train_loop - len(code) - 1) & 0xFF)  # displacement from PC after DJNZ
    emit(0xAF)                     # XOR A           ; A = 0 for bank select
    emit(0xC9)                     # RET

    # =====================================================================
    # Data area
    # =====================================================================
    symbols['RFSCR'] = addr()
    emit(0x0D)                     # DB 0DH  ; CR = empty filename (next file)

    symbols['RFS_FIRST'] = addr()
    emit(0x00)                     # DB 00H  ; First-call flag for ZRO

    symbols['SAVE_SP'] = addr()
    emit(0x00, 0x00)              # DW 0000H ; Saved SP during ROM calls

    # Temp stack: 32 bytes (16 levels) for ROM calls while upper area is MMIO
    symbols['TEMP_STACK_BASE'] = addr()
    code.extend(b'\x00' * 32)
    symbols['TEMP_STACK_TOP'] = addr()  # SP points here (grows downward)

    # =====================================================================
    # Patch all forward references
    # =====================================================================
    rfs_first_addr = symbols['RFS_FIRST']
    save_sp_addr = symbols['SAVE_SP']
    temp_stack_top = symbols['TEMP_STACK_TOP']

    # RFSINIT: LD (RFS_FIRST),A
    struct.pack_into('<H', code, rfsinit_first_patch + 1, rfs_first_addr)

    # RFSRO: LD A,(RFS_FIRST)
    struct.pack_into('<H', code, rfsro_first_load + 1, rfs_first_addr)

    # RFSRO: LD (RFS_FIRST),A
    struct.pack_into('<H', code, rfsro_first_store + 1, rfs_first_addr)

    # RFSRO: LD DE,RFSCR
    struct.pack_into('<H', code, rfscr_load + 1, symbols['RFSCR'])

    # RFSRO: CALL CALL_RDINF
    struct.pack_into('<H', code, rfsro_call_rdinf, symbols['CALL_RDINF'])

    # RFSWO: CALL CALL_WRINF
    struct.pack_into('<H', code, rfswo_call_wrinf, symbols['CALL_WRINF'])

    # RFSINP: CALL CALL_RDDATA
    struct.pack_into('<H', code, rfsinp_call_rddata, symbols['CALL_RDDATA'])

    # RFSOUT: CALL CALL_WRDATA
    struct.pack_into('<H', code, rfsout_call_wrdata, symbols['CALL_WRDATA'])

    # Patch all 4 CALL TRAIN_UPCNTR forward references
    train_addr = symbols['TRAIN_UPCNTR']
    struct.pack_into('<H', code, call_rdinf_train + 1,  train_addr)
    struct.pack_into('<H', code, call_rddata_train + 1, train_addr)
    struct.pack_into('<H', code, call_wrinf_train + 1,  train_addr)
    struct.pack_into('<H', code, call_wrdata_train + 1, train_addr)

    # Patch all 4 ROM call wrappers: SAVE_SP and TEMP_STACK_TOP addresses
    for save_off, set_off, rest_off in [
        (call_rdinf_save_sp,  call_rdinf_set_sp,  call_rdinf_rest_sp),
        (call_rddata_save_sp, call_rddata_set_sp,  call_rddata_rest_sp),
        (call_wrinf_save_sp,  call_wrinf_set_sp,   call_wrinf_rest_sp),
        (call_wrdata_save_sp, call_wrdata_set_sp,   call_wrdata_rest_sp),
    ]:
        struct.pack_into('<H', code, save_off + 2, save_sp_addr)     # LD (SAVE_SP),SP
        struct.pack_into('<H', code, set_off + 1,  temp_stack_top)   # LD SP,TEMP_STACK_TOP
        struct.pack_into('<H', code, rest_off + 2, save_sp_addr)     # LD SP,(SAVE_SP)

    # =====================================================================
    # Fill in the device descriptor (29 bytes at offset 0)
    # =====================================================================
    desc = bytearray(29)
    struct.pack_into('<H', desc, 0, USR_ADDR)          # Next = $USR (0D5EH)
    desc[2] = ord('R')
    desc[3] = ord('F')
    desc[4] = ord('S')
    desc[5] = 0x00                                      # Name null terminator
    desc[6] = 0x5F                                      # Flags1: SEQ, FNM, W, R
    desc[7] = 0x07                                      # Flags2: max channel 7 (bit7 must be 0; SVC#39 errors on bit7)
    desc[8] = 0xFF                                      # Dirmx: 255 entries
    struct.pack_into('<H', desc, 9,  symbols['RFSINIT'])  # ZINIT
    struct.pack_into('<H', desc, 11, symbols['RFSRO'])    # ZRO
    struct.pack_into('<H', desc, 13, symbols['RFSWO'])    # ZWO
    struct.pack_into('<H', desc, 15, symbols['RFSSTRT'])  # ZCL/ZSTRT
    struct.pack_into('<H', desc, 17, 0x0000)              # ZKL (not used)
    struct.pack_into('<H', desc, 19, symbols['RFSINP'])   # ZINP
    struct.pack_into('<H', desc, 21, symbols['RFSOUT'])   # ZOUT
    struct.pack_into('<H', desc, 23, 0x0000)              # ZDELT (not used)
    struct.pack_into('<H', desc, 25, 0x0000)              # ZWDIR (not used)
    struct.pack_into('<H', desc, 27, symbols['RFSFREE'])  # ZFREE
    code[desc_offset:desc_offset + 29] = desc

    return bytes(code), symbols


def patch_binary(binary_data, rfs_code, symbols):
    """
    Patch the MZ-5Z008 binary to add RFS device support.

    Modifications:
    1. Append RFS device code at 85FFH (end of binary)
    2. Patch $QD link at 38A2H from 0D5EH ($USR) to 85FFH ($RFS)
    3. Update TEXTST init at 4000H to point past RFS code
    4. Update MEMCLI start at 4020H: RST 18H/DB 20H/EX DE,HL -> LD HL,new_textst

    Returns the patched binary.
    """
    if len(binary_data) != BINARY_SIZE:
        raise ValueError(
            f"Binary must be exactly {BINARY_SIZE} bytes, got {len(binary_data)}"
        )

    data = bytearray(binary_data)
    new_textst = RFS_BASE + len(rfs_code)

    # Verify $QD link pointer is currently 0D5EH ($USR)
    qd_link = struct.unpack_from('<H', data, QD_LINK)[0]
    if qd_link != USR_ADDR:
        print(f"WARNING: $QD link at {QD_LINK:04X}H is {qd_link:04X}H, expected {USR_ADDR:04X}H")
        print("         Binary may already be patched or is not the expected version.")

    # Verify TEXTST init at 4000H: should be "21 FF 85" (LD HL,85FFH)
    if data[TEXTST_INIT] != 0x21 or data[TEXTST_INIT+1] != 0xFF or data[TEXTST_INIT+2] != 0x85:
        print(f"WARNING: Expected LD HL,85FFH (21 FF 85) at {TEXTST_INIT:04X}H, "
              f"got {data[TEXTST_INIT]:02X} {data[TEXTST_INIT+1]:02X} {data[TEXTST_INIT+2]:02X}")

    # Verify MEMCLI init at 4020H: should be "DF 20 EB" (RST 18H / DB 20H / EX DE,HL)
    if data[MEMCLI_INIT] != 0xDF or data[MEMCLI_INIT+1] != 0x20 or data[MEMCLI_INIT+2] != 0xEB:
        print(f"WARNING: Expected DF 20 EB at {MEMCLI_INIT:04X}H, "
              f"got {data[MEMCLI_INIT]:02X} {data[MEMCLI_INIT+1]:02X} {data[MEMCLI_INIT+2]:02X}")

    # Patch 1: Append RFS device code
    data.extend(rfs_code)
    print(f"  Appended {len(rfs_code)} bytes of RFS code at {RFS_BASE:04X}H-{RFS_BASE + len(rfs_code) - 1:04X}H")

    # Patch 2: Update $QD link pointer: $USR -> $RFS
    struct.pack_into('<H', data, QD_LINK, RFS_BASE)
    print(f"  Patched $QD link at {QD_LINK:04X}H: {USR_ADDR:04X}H -> {RFS_BASE:04X}H")

    # Patch 3: Update TEXTST init: LD HL,85FFH -> LD HL,new_textst
    data[TEXTST_INIT + 1] = new_textst & 0xFF
    data[TEXTST_INIT + 2] = (new_textst >> 8) & 0xFF
    print(f"  Patched TEXTST init at {TEXTST_INIT:04X}H: 85FFH -> {new_textst:04X}H")

    # Patch 4: Update MEMCLI start: RST 18H/DB 20H/EX DE,HL -> LD HL,new_textst
    # RST 18H (DF) + DB 20H + EX DE,HL (EB) = 3 bytes -> LD HL,nn (21 lo hi) = 3 bytes
    data[MEMCLI_INIT]     = 0x21
    data[MEMCLI_INIT + 1] = new_textst & 0xFF
    data[MEMCLI_INIT + 2] = (new_textst >> 8) & 0xFF
    print(f"  Patched MEMCLI start at {MEMCLI_INIT:04X}H: DF 20 EB -> 21 {new_textst & 0xFF:02X} {(new_textst >> 8) & 0xFF:02X}")

    # Patch 5: Bump version letter in "V1.0A" string
    old_ver = chr(data[VERSION_OFFSET])
    data[VERSION_OFFSET] = VERSION_LETTER
    print(f"  Patched version at {VERSION_OFFSET:04X}H: '{old_ver}' -> '{chr(VERSION_LETTER)}'")

    print(f"  Device chain: $QD -> $RFS ({RFS_BASE:04X}H) -> $USR ({USR_ADDR:04X}H) -> NULL")
    print(f"  BASIC program text area: {new_textst:04X}H-FEFFH "
          f"({0xFEFF - new_textst} bytes, was {0xFEFF - 0x85FF} bytes, "
          f"reduced by {len(rfs_code)} bytes)")

    return bytes(data)


def main():
    input_bin = sys.argv[1] if len(sys.argv) > 1 else 'mz5z008.bin'
    output_bin = sys.argv[2] if len(sys.argv) > 2 else None

    if output_bin is None:
        stem = os.path.splitext(input_bin)[0]
        output_bin = stem + '_rfs.bin'

    # Build RFS code
    rfs_code, symbols = build_rfs_code()
    new_textst = RFS_BASE + len(rfs_code)
    print(f"RFS device module: {len(rfs_code)} bytes at {RFS_BASE:04X}H-{RFS_BASE + len(rfs_code) - 1:04X}H")
    print(f"  New TEXTST: {new_textst:04X}H (was 85FFH, {len(rfs_code)} bytes less for BASIC programs)")
    print()

    # Print symbol table
    print("Symbol table:")
    for name, addr in sorted(symbols.items(), key=lambda x: x[1]):
        print(f"  {name:12s} = {addr:04X}H")
    print()

    # Read original binary
    with open(input_bin, 'rb') as f:
        binary_data = f.read()

    if len(binary_data) != BINARY_SIZE:
        print(f"ERROR: Binary must be exactly {BINARY_SIZE} bytes, got {len(binary_data)}")
        sys.exit(1)

    # Apply patches
    print("Patching binary:")
    patched = patch_binary(binary_data, rfs_code, symbols)

    # Write output
    with open(output_bin, 'wb') as f:
        f.write(patched)
    print(f"\nWrote {len(patched)} bytes to {output_bin}")

    # Verify the patch
    print("\nVerification:")
    # Check $QD link
    qd_link = struct.unpack_from('<H', patched, QD_LINK)[0]
    print(f"  $QD link at {QD_LINK:04X}H = {qd_link:04X}H {'OK' if qd_link == RFS_BASE else 'FAIL'}")
    # Check $RFS descriptor
    rfs_next = struct.unpack_from('<H', patched, RFS_BASE)[0]
    rfs_name = patched[RFS_BASE + 2:RFS_BASE + 6]
    rfs_flags1 = patched[RFS_BASE + 6]
    print(f"  $RFS next at {RFS_BASE:04X}H = {rfs_next:04X}H {'OK' if rfs_next == USR_ADDR else 'FAIL'}")
    expected_name = b'RFS\x00'
    print(f"  $RFS name = {rfs_name} {'OK' if rfs_name == expected_name else 'FAIL'}")
    print(f"  $RFS flags1 = {rfs_flags1:02X}H {'OK' if rfs_flags1 == 0x5F else 'FAIL'}")
    # Check TEXTST
    textst_val = patched[TEXTST_INIT+1] | (patched[TEXTST_INIT+2] << 8)
    textst_ok = textst_val == new_textst
    print(f"  TEXTST at {TEXTST_INIT:04X}H = {textst_val:04X}H {'OK' if textst_ok else 'FAIL'}")
    # Check MEMCLI
    memcli_val = patched[MEMCLI_INIT+1] | (patched[MEMCLI_INIT+2] << 8)
    memcli_ok = memcli_val == new_textst
    print(f"  MEMCLI at {MEMCLI_INIT:04X}H = {memcli_val:04X}H {'OK' if memcli_ok else 'FAIL'}")
    # Check MEMCLI opcode changed from RST to LD HL
    memcli_op = patched[MEMCLI_INIT]
    print(f"  MEMCLI opcode at {MEMCLI_INIT:04X}H = {memcli_op:02X}H {'OK' if memcli_op == 0x21 else 'FAIL'}")
    # Check device chain integrity
    print(f"\n  Device chain walk:")
    off = QD_LINK
    link = struct.unpack_from('<H', patched, off)[0]
    print(f"    $QD at {QD_LINK:04X}H -> link={link:04X}H")
    if link == RFS_BASE:
        rfs_off = RFS_BASE
        rfs_link = struct.unpack_from('<H', patched, rfs_off)[0]
        rfs_name_str = patched[rfs_off+2:rfs_off+6].rstrip(b'\x00').decode()
        print(f"    ${rfs_name_str} at {rfs_off:04X}H -> link={rfs_link:04X}H")
        if rfs_link == USR_ADDR:
            usr_link = struct.unpack_from('<H', patched, USR_ADDR)[0]
            usr_name_str = patched[USR_ADDR+2:USR_ADDR+6].rstrip(b'\x00').decode()
            print(f"    ${usr_name_str} at {USR_ADDR:04X}H -> link={usr_link:04X}H (NULL=end)")

    return output_bin


if __name__ == '__main__':
    result = main()
