"""Locate the host-side mirror of PSP MMIO 0x03010000 via the crash-dump magic.

The PSP fatal handler writes:
    MEMORY[0x3010A48] = 0
    MEMORY[0x3010A4C] = (exc_class << 28) | (x << 24) | 0xE504FF   <-- magic
    MEMORY[0x3010A50] = faulting LR
    MEMORY[0x3010A54] = DFSR/IFSR
    MEMORY[0x3010A58] = DFAR/IFAR

Plan: snapshot all of BAR5, deliberately crash the PSP, snapshot again,
then hunt for a dword whose low 24 bits are 0xE504FF.
"""
import ctypes, struct, time, array
from tinygrad import Device
from tinygrad.runtime.autogen.am import am

d = Device["AMD"]; adev = d.iface.dev_impl; psp = adev.psp
SZ = ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
NDW = len(adev.mmio)
CHUNK = 0x4000

def scan():
    out = array.array('I')
    for base in range(0, NDW, CHUNK):
        n = min(CHUNK, NDW - base)
        out.extend(array.array('I', bytes(adev.mmio[base:base+n])[:n*4]))
    return out

print(f"sOS alive: {psp.is_sos_alive()}   scanning {NDW} dwords ({NDW*4/1e6:.2f} MB)")
t0 = time.perf_counter(); before = scan()
print(f"  baseline scan: {time.perf_counter()-t0:.2f}s, {len(before)} dwords")

MAGIC = 0xE504FF
pre_hits = [i for i, v in enumerate(before) if (v & 0xFFFFFF) == MAGIC]
print(f"  magic {MAGIC:#08x} present BEFORE crash: {len(pre_hits)} {[hex(x) for x in pre_hits[:5]]}")

print("\n>>> crashing the PSP (LOAD_IP_FW fw_size=0x100, addr=0) ...")
adev.vram.view(psp.cmd_paddr, SZ)[:] = b"\x00" * SZ
c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_IP_FW)
c.cmd.cmd_load_ip_fw.fw_size = 0x100
c.cmd.cmd_load_ip_fw.fw_type = 0
try:
    psp._ring_submit(c); print("  (unexpectedly returned)")
except Exception as e:
    print(f"  {type(e).__name__}: {str(e)[:70]}")
time.sleep(1.0)

t0 = time.perf_counter(); after = scan()
print(f"  post-crash scan: {time.perf_counter()-t0:.2f}s")

hits = [i for i, v in enumerate(after) if (v & 0xFFFFFF) == MAGIC]
print(f"\n=== dwords with low-24 == {MAGIC:#08x} AFTER crash: {len(hits)} ===")
for i in hits:
    print(f"  dword {i:#08x} (byte {i*4:#08x}) = {after[i]:#010x}   exc_class={(after[i]>>28)&0xF}")
    base_dw = i - (0xA4C // 4)
    print(f"     => PSP 0x03010000 maps to BAR5 dword {base_dw:#08x} (byte {base_dw*4:#08x})")
    print(f"     neighbourhood (PSP 0x3010A48..0x3010A60):")
    for off in (0xA48, 0xA4C, 0xA50, 0xA54, 0xA58, 0xA5C, 0xA60):
        j = base_dw + off//4
        if 0 <= j < len(after):
            print(f"       0x{0x3010000+off:08x} -> dw {j:#08x}  before={before[j]:#010x}  after={after[j]:#010x}")

diff = [(i, before[i], after[i]) for i in range(len(before)) if before[i] != after[i]]
print(f"\n=== total dwords changed by the crash: {len(diff)} ===")
for i, b, a in diff[:40]:
    print(f"  dw {i:#08x} (byte {i*4:#08x})  {b:#010x} -> {a:#010x}")
if len(diff) > 40: print(f"  ... and {len(diff)-40} more")