#!/usr/bin/env python3
"""
Patch ESP32 QEMU source to fix WiFi/RF peripheral register stubs.

Two changes only — no new files, no meson.build modifications:

  1. hw/misc/unimp.c — change unimp_read to return 0xFFFFFFFF instead of 0.
     The ESP32 WiFi library busy-loops on hardware-ready bits.  With the
     default return-0 those loops never exit and the watchdog fires.
     0xFFFF... sets all bits, so most active-high "ready/done" flags pass.

  2. hw/xtensa/esp32.c — add create_unimplemented_device() calls for all
     unmapped WiFi/RF peripheral registers that cause LoadStorePIFAddrError.
"""
import re
import sys

ESP32_C = 'hw/xtensa/esp32.c'
UNIMP_C = 'hw/misc/unimp.c'

# ── Stubs to register ─────────────────────────────────────────────────────────
STUB_PREFIX = 'rfstub.'   # unimp_read returns 0xFF for names starting with this

STUBS = [
    # (device-name, base-addr, size)
    # Prefix "rfstub." distinguishes our WiFi stubs from other unimplemented
    # devices — unimp_read returns 0xFF only for that prefix.

    # AHB bus ─────────────────────────────────────────────────────────────────
    (f'{STUB_PREFIX}wifi_modem',  0x60033C00, 0x10000),

    # Data-bus (0x3FF00000) ───────────────────────────────────────────────────
    (f'{STUB_PREFIX}fe2',         0x3FF45000, 0x1000),
    (f'{STUB_PREFIX}fe',          0x3FF46000, 0x1000),   # crash @ 0x3FF460A0
    (f'{STUB_PREFIX}bt_bb',       0x3FF51000, 0x2000),   # <=0x2000 avoids I2C_EXT0
    (f'{STUB_PREFIX}5c',          0x3FF5C000, 0x1000),   # crash @ 0x3FF5C01C
    (f'{STUB_PREFIX}5d',          0x3FF5D000, 0x1000),
    (f'{STUB_PREFIX}nrx',         0x3FF5E000, 0x1000),   # 0x3FF5F000=TG0 untouched
    (f'{STUB_PREFIX}62',          0x3FF62000, 0x1000),
    (f'{STUB_PREFIX}63',          0x3FF63000, 0x1000),
    (f'{STUB_PREFIX}68',          0x3FF68000, 0x1000),
    (f'{STUB_PREFIX}6a',          0x3FF6A000, 0x1000),
    (f'{STUB_PREFIX}6b',          0x3FF6B000, 0x1000),
    (f'{STUB_PREFIX}6c',          0x3FF6C000, 0x1000),
    # 0x3FF6E000 = UART2 — leave for Modbus
]


# ─────────────────────────────────────────────────────────────────────────────
# Step 1 — patch unimp_read to return 0xFFFFFFFF
# ─────────────────────────────────────────────────────────────────────────────

def patch_unimp(path):
    try:
        with open(path) as f:
            src = f.read()
    except FileNotFoundError:
        print(f'  ERROR: {path} not found', file=sys.stderr)
        sys.exit(1)

    if '~(uint64_t)0' in src or '0xffffffff' in src.lower() and 'unimp_read' in src:
        print(f'  skip {path} — already patched')
        return

    # Replace "return 0;" inside unimp_read with a prefix-checked return:
    # devices named "rfstub.*" return 0xFFFFFFFF, all others return 0.
    replacement = (
        'return strncmp(s->name, "rfstub.", 7) == 0 ? ~(uint64_t)0 : 0;'
    )
    new_src = re.sub(
        r'(static uint64_t unimp_read\b[^}]+?\breturn\s+)0(;)',
        lambda m: m.group(0).replace('return 0;', replacement),
        src,
        count=1,
        flags=re.DOTALL,
    )
    # Also ensure <string.h> is included for strncmp
    if '#include <string.h>' not in new_src and 'strncmp' in new_src:
        new_src = new_src.replace(
            '#include "qemu/osdep.h"\n',
            '#include "qemu/osdep.h"\n#include <string.h>\n',
            1,
        )
    if new_src == src:
        print(f'  WARNING: pattern not found in {path}, skipping', file=sys.stderr)
        return

    with open(path, 'w') as f:
        f.write(new_src)
    print(f'  patched {path}: unimp_read returns 0xFFFF for rfstub.* only')


# ─────────────────────────────────────────────────────────────────────────────
# Step 2 — register stubs in ESP32 machine
# ─────────────────────────────────────────────────────────────────────────────

def patch_esp32(path):
    try:
        with open(path) as f:
            src = f.read()
    except FileNotFoundError:
        print(f'  ERROR: {path} not found', file=sys.stderr)
        sys.exit(1)

    # Ensure hw/misc/unimp.h is included
    if 'hw/misc/unimp.h' not in src:
        src = src.replace(
            '#include "qemu/osdep.h"\n',
            '#include "qemu/osdep.h"\n#include "hw/misc/unimp.h"\n',
            1,
        )
        print('  + added #include "hw/misc/unimp.h"')

    changed = False
    for name, addr, size in STUBS:
        hex_addr = f'0x{addr:08X}'
        if hex_addr in src or hex_addr.lower() in src:
            print(f'  skip {name}: already present')
            continue

        stub_line = (
            f'    create_unimplemented_device("{name}", {hex_addr}, 0x{size:X});\n'
        )

        m = list(re.finditer(r'create_unimplemented_device\([^;]+;\n', src))
        if m:
            pos = m[-1].end()
            src = src[:pos] + stub_line + src[pos:]
        else:
            pos = src.rfind('\n}')
            src = src[:pos] + '\n' + stub_line + src[pos:]
        print(f'  + stub {name} @ {hex_addr}')
        changed = True

    if changed:
        with open(path, 'w') as f:
            f.write(src)
        print(f'  wrote {path}')
    else:
        print(f'  no changes to {path}')


# ─────────────────────────────────────────────────────────────────────────────

if __name__ == '__main__':
    print('Step 1: patch unimp_read to return 0xFFFFFFFF')
    patch_unimp(UNIMP_C)

    print('Step 2: register WiFi/RF stubs in ESP32 machine')
    patch_esp32(ESP32_C)

    print('Patch complete.')