import ctypes
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); C=28
def run(cid,words):
    adev.vram.view(psp.cmd_paddr,SZ)[:]=b"\x00"*SZ
    c=am.struct_psp_gfx_cmd_resp(cmd_id=cid); buf=bytearray(memoryview(c).cast('B'))
    for i,v in words.items(): buf[C+i*4:C+4+i*4]=int(v&0xFFFFFFFF).to_bytes(4,'little')
    adev.vram.view(psp.cmd_paddr,SZ)[:]=bytes(buf)
    before=bytes(buf)
    try: psp._ring_submit(am.struct_psp_gfx_cmd_resp.from_buffer_copy(bytes(buf)))
    except Exception: pass
    after=bytes(adev.vram.view(psp.cmd_paddr,SZ)[:])
    return before,after
for label,words in (("READ idx=5",{0:5,1:0,2:1}),("WRITE idx=5 magic",{0:5,1:0xA5A5F00D,2:0}),("READ idx=3",{0:3,1:0,2:1})):
    b,a=run(0x0d,words)
    diffs=[(o,int.from_bytes(b[o:o+4],'little'),int.from_bytes(a[o:o+4],'little')) for o in range(0,SZ,4)
           if b[o:o+4]!=a[o:o+4]]
    print(f"\ncmd 0x0d {label}: {len(diffs)} words changed by the PSP")
    for o,x,y in diffs: print(f"   +{o:#05x} {x:#010x} -> {y:#010x}")
print("\nsOS alive:",psp.is_sos_alive())