"""Provoke distinctive error codes from the 5 undocumented commands that
currently return 0 with null parameters.

Method (same as PROG_REG): vary one dword of the command union at a time.
Any status != 0 is a constant we can pivot on to locate the handler.
"""
import ctypes, time, json
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)
CMDS = [0x0c, 0x0d, 0x12, 0x24, 0x30]
VALUES = [1, 2, 0x10, 0xFF, 0x100, 0xFFFFFFFF]
NWORDS = 6

def alive(): return adev.reg(f"{psp.reg_pref}_81").read() != 0

def submit(cmd_id, words):
    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[28 + i*4 : 32 + i*4] = int(v & 0xFFFFFFFF).to_bytes(4, 'little')
    try:
        psp._ring_submit(am.struct_psp_gfx_cmd_resp.from_buffer_copy(bytes(buf)))
        return 0
    except RuntimeError:
        r = bytes(adev.vram.view(psp.cmd_paddr, SZ)[:])
        return int.from_bytes(r[864:868], "little")
    except TimeoutError:
        return -1

results = {}
statuses = {}
print(f"{'cmd':>5s} {'word':>5s} {'value':>12s} {'status':>12s}  note")
aborted = False
for cid in CMDS:
    if aborted: break
    base = submit(cid, {})
    print(f"{cid:#05x} {'-':>5s} {'(all zero)':>12s} {base if base<0 else hex(base):>12s}  baseline")
    for wi in range(NWORDS):
        if aborted: break
        for v in VALUES:
            if not alive():
                print("!! sOS died -- aborting"); aborted = True; break
            s = submit(cid, {wi: v})
            results[f"{cid:#x}:{wi}:{v:#x}"] = s
            if s != base:
                ss = "TIMEOUT" if s < 0 else hex(s)
                statuses.setdefault(s, []).append((cid, wi, v))
                print(f"{cid:#05x} {wi:>5d} {v:#012x} {ss:>12s}  <-- DIFFERS from baseline")
    json.dump(results, open("/tmp/undoc_probe.json", "w"))

print(f"\nsOS alive: {alive()}")
print(f"\n=== distinct non-baseline statuses ===")
for s, hits in sorted(statuses.items()):
    ss = "TIMEOUT" if s < 0 else hex(s)
    print(f"  {ss:>12s} : {len(hits)} hits, e.g. cmd={hits[0][0]:#x} word={hits[0][1]} val={hits[0][2]:#x}")
if not statuses:
    print("  none -- all commands ignore these parameters")