"""Test whether undocumented cmd 0x0c / 0x0d expose the sub_E13A2C read/write
primitive over MMIO window 0x03010500 + idx*4.

Predicted layout (from decompilation):
    word0 = index (bounded)
    word1 = value in/out
    word2 = direction: 0 -> WRITE, nonzero -> READ

Decisive test: READ an index, WRITE a different value, READ back.
If readback reflects the write, the primitive is host-controlled.
We restore the original value afterwards.
"""
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)
CMDOFF = 28

def submit(cmd_id, words):
    """returns (status, full command buffer after the PSP processed it)"""
    adev.vram.view(psp.cmd_paddr, SZ)[:] = b"\x00" * SZ
    c = am.struct_psp_gfx_cmd_resp(cmd_id=cmd_id)
    buf = bytearray(memoryview(c).cast('B'))
    for i, v in words.items():
        buf[CMDOFF+i*4:CMDOFF+4+i*4] = int(v & 0xFFFFFFFF).to_bytes(4, 'little')
    adev.vram.view(psp.cmd_paddr, SZ)[:] = bytes(buf)
    st = 0
    try:
        psp._ring_submit(am.struct_psp_gfx_cmd_resp.from_buffer_copy(bytes(buf)))
    except RuntimeError:
        pass
    except TimeoutError:
        st = -1
    out = bytes(adev.vram.view(psp.cmd_paddr, SZ)[:])
    if st == 0:
        st = int.from_bytes(out[864:868], "little")
    return st, out

def word(buf, i): return int.from_bytes(buf[CMDOFF+i*4:CMDOFF+4+i*4], "little")

for cid in (0x0c, 0x0d):
    print(f"\n================ cmd {cid:#04x} ================")
    IDX = 5
    st, out = submit(cid, {0: IDX, 1: 0, 2: 1})        # READ
    orig = word(out, 1)
    print(f"  READ  idx={IDX}          -> status={hex(st):>12s}  word1={orig:#010x}")

    MAGIC = 0xA5A5F00D
    st2, out2 = submit(cid, {0: IDX, 1: MAGIC, 2: 0})  # WRITE
    print(f"  WRITE idx={IDX} val={MAGIC:#010x} -> status={hex(st2):>12s}")

    st3, out3 = submit(cid, {0: IDX, 1: 0, 2: 1})      # READ BACK
    back = word(out3, 1)
    print(f"  READ  idx={IDX}          -> status={hex(st3):>12s}  word1={back:#010x}")

    if back == MAGIC:
        print("  *** READBACK MATCHES THE WRITE -- host-controlled R/W primitive ***")
    elif back != orig:
        print(f"  *** value CHANGED ({orig:#x} -> {back:#x}) but != magic ***")
    else:
        print("  no observable change (word1 not returned, or writes ignored)")

    if back != orig:
        submit(cid, {0: IDX, 1: orig, 2: 0})           # restore
        _, chk = submit(cid, {0: IDX, 1: 0, 2: 1})
        print(f"  restored -> {word(chk,1):#010x} (was {orig:#010x})")

print("\nsOS alive:", psp.is_sos_alive())