#!/usr/bin/env python3
"""
patch_rfs_mz5z009.py - Patch MZ-5Z009 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 NEWAD0 (A471H), the current end of the binary.
The $QD device's link pointer is patched to chain: $QD → $RFS → $USR.
TEXTST and MEMCLI are moved past the RFS code.

Unlike MZ-2Z009E, no DIR hook is needed — MZ-5Z009 uses the standard SVC .DIR
mechanism which calls FDIR in the monitor. FDIR calls ZRO in a loop to enumerate
files for SEQ devices, so the RFS device's ZRO handler is called automatically.

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

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

  input_binary:  Original MZ-5Z009 binary (default: mz5z009.bin)
  output_binary: Patched binary with RFS support (default: mz5z009_rfs.bin)
"""

import struct
import sys
import os

# --- Constants ---
BINARY_SIZE     = 42097              # Expected MZ-5Z009 binary size (0xA471 bytes)
RFS_BASE        = 0xA471             # Append RFS code at end of binary (= NEWAD0)
QD_LINK         = 0x2FD0             # $QD device descriptor link pointer
USR_ADDR        = 0x0F7F             # $USR device descriptor address
TEXTST_INIT     = 0x5805             # "LD HL,A471H" in BASIC init
MEMCLI_INIT     = 0x5837             # "LD HL,A471H" for MEMCLI clearing

# MZ-5Z009 addresses
ELMD            = 0x1000
ELMD1           = 0x1001
IBUFE           = 0x10F0
ZCH             = 0x1046
SDDRIVENO       = 0x1028
CMTFILENO        = 0x1029
BREAKZ          = 0x00E8
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.
    """
    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))

    # =====================================================================
    # $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

    # =====================================================================
    # 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 (ELMD filled), CY = no more files
    # =====================================================================
    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 empty filename (CR byte) to get next file
    rfscr_load = len(code)
    emit(0x11, 0x00, 0x00)        # LD DE,RFSCR       [patched below]
    emit(0xCD, *le16(RFSF_RDINF)) # CALL RFSF_RDINF
    rfsro_jrc = len(code)
    emit(0x38, 0x00)              # JR C,RFSROERR     [patched below]
    # Copy IBUFE -> ELMD (header data)
    emit(0x21, *le16(IBUFE))      # LD HL,IBUFE
    emit(0x11, *le16(ELMD))       # LD DE,ELMD
    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
    emit(0x11, *le16(IBUFE + 1))  # LD DE,IBUFE+1 (NAME field)
    emit(0xCD, *le16(RFSF_WRINF)) # CALL RFSF_WRINF
    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(RFSF_RDDATA))  # CALL RFSF_RDDATA
    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(RFSF_WRDATA))  # CALL RFSF_WRDATA
    emit(0xC1)                     # POP BC
    emit(0xD1)                     # POP DE
    emit(0xE1)                     # POP HL
    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

    # =====================================================================
    # Patch all forward references
    # =====================================================================
    rfs_first_addr = symbols['RFS_FIRST']

    # 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'])

    # =====================================================================
    # Fill in the device descriptor (29 bytes at offset 0)
    # =====================================================================
    desc = bytearray(29)
    struct.pack_into('<H', desc, 0, USR_ADDR)          # Next = $USR (0F7FH)
    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
    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-5Z009 binary to add RFS device support.

    Modifications:
    1. Append RFS device code at A471H (end of binary)
    2. Patch $QD link at 2FD0H from 0F7FH ($USR) to A471H ($RFS)
    3. Update TEXTST init at 5805H to point past RFS code
    4. Update MEMCLI start at 5837H to point past RFS code

    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 0F7FH ($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 5805H: should be "21 71 A4" (LD HL,A471H)
    if data[TEXTST_INIT] != 0x21 or data[TEXTST_INIT+1] != 0x71 or data[TEXTST_INIT+2] != 0xA4:
        print(f"WARNING: Expected LD HL,A471H (21 71 A4) 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 5837H: should be "21 71 A4" (LD HL,A471H)
    if data[MEMCLI_INIT] != 0x21 or data[MEMCLI_INIT+1] != 0x71 or data[MEMCLI_INIT+2] != 0xA4:
        print(f"WARNING: Expected LD HL,A471H (21 71 A4) 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
    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: A471H -> {new_textst:04X}H")

    # Patch 4: Update MEMCLI start
    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: A471H -> {new_textst:04X}H")

    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 - 0xA471} bytes, "
          f"reduced by {len(rfs_code)} bytes)")

    return bytes(data)


def main():
    input_bin = sys.argv[1] if len(sys.argv) > 1 else 'mz5z009.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 A471H, {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 device chain integrity
    print(f"\n  Device chain walk:")
    # Start from $QD
    off = QD_LINK
    name = patched[off+2:off+6]
    link = struct.unpack_from('<H', patched, off)[0]
    print(f"    $QD at {QD_LINK:04X}H (in binary) -> link={link:04X}H")
    # Follow to $RFS
    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")
        # Follow to $USR
        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()
