#!/usr/bin/env python3
"""Sweep the PSP ring CONTROL command space written to C2PMSG_64.

Distinct from the ring command space: `_64` takes a control id in bits 16-19, so there are 16
possible commands, of which eleven are named in the headers (CAN_INIT_RINGS, INIT_RBI_RING,
INIT_GPCOM_RING, DESTROY_RINGS, DESTROY_GPCOM_RING, ENABLE_INT, DISABLE_INT, GBR_IH_SET,
CONSUME_CMD, MODE1_RST, MAX) and five are undocumented. Only ring-create and DESTROY_RINGS
have ever been used here.

Per id: the control runs FIRST in the same boot (ring creates, a real command completes), then
the id is written and the status read, then the ring is recreated and a command submitted
again. An id that leaves the PSP unable to run commands is a denial of service.

MODE1_RST (0x70000) resets the ASIC by design -- not a finding, and it costs a recovery
cycle -- so it is not in the caller's default list.

argv[1]: hex id, or 'list'
"""
import ctypes, glob, struct, sys, time
sys.path.insert(0, '/root/psp/hw')
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am
adev = Device["AMD"].iface.dev_impl; psp = adev.psp
RESP = 864; FRAME = ctypes.sizeof(am.struct_psp_gfx_rb_frame); step = FRAME // 4
P = psp.reg_pref
NAMES = {getattr(am, n): n.replace('GFX_CTRL_CMD_ID_', '')
         for n in dir(am) if 'GFX_CTRL_CMD_ID' in n}
def alloc(n):
    pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False); return pa, adev.paddr2mc(pa)
MAGIC = 0x6C710001
blob = open(glob.glob('/root/psp/contall/ta_42-0-0-6_1c874c3c*')[0], 'rb').read()
ap_, amc = alloc(0x80000); cbuf, cbmc = alloc(0x8000); scratch, smc = alloc(0x1000)
n = (len(blob) + 3) & ~3
adev.vram.view(ap_, n)[:] = blob + b'\x00' * (n - len(blob)); adev.gmc.flush_hdp()
ring_mc = adev.paddr2mc(psp.ring_paddr)
def create():
    if adev.reg(f"{P}_71").read() != 0:
        adev.reg(f"{P}_64").write(am.GFX_CTRL_CMD_ID_DESTROY_RINGS); time.sleep(0.05)
    t0 = time.time()
    while time.time() - t0 < 3:
        if adev.reg(f"{P}_64").read() & 0x80000000: break
    adev.wreg_pair(P, "_69", "_70", ring_mc)
    adev.reg(f"{P}_71").write(psp.ring_size)
    adev.reg(f"{P}_64").write(am.PSP_RING_TYPE__KM << 16)
    time.sleep(0.1)
    return (adev.reg(f"{P}_64").read() & 0x8000FFFF) == 0x80000000
def submit():
    c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA); t = c.cmd.cmd_load_ta
    t.app_phy_addr_lo, t.app_phy_addr_hi = amc & 0xFFFFFFFF, amc >> 32; t.app_len = len(blob)
    t.cmd_buf_phy_addr_lo, t.cmd_buf_phy_addr_hi = cbmc & 0xFFFFFFFF, cbmc >> 32
    t.cmd_buf_len = 0x4000
    adev.vram.view(psp.cmd_paddr, ctypes.sizeof(c))[:] = memoryview(c).cast('B')
    adev.vram.view(psp.cmd_paddr + RESP, 4, 'I')[0] = 0xA5A5A5A5
    adev.vram.view(scratch, 4, 'I')[0] = 0
    f = am.struct_psp_gfx_rb_frame(fence_value=MAGIC,
        cmd_buf_addr_lo=adev.paddr2mc(psp.cmd_paddr) & 0xFFFFFFFF,
        cmd_buf_addr_hi=adev.paddr2mc(psp.cmd_paddr) >> 32,
        fence_addr_lo=smc & 0xFFFFFFFF, fence_addr_hi=smc >> 32)
    adev.vram.view(psp.ring_paddr, FRAME)[:] = bytes(memoryview(f).cast('B'))
    adev.gmc.flush_hdp()
    adev.reg(f"{P}_67").write(step)
    t0 = time.time()
    while time.time() - t0 < 2.0:
        adev.gmc.flush_hdp()
        v = struct.unpack_from('<I', bytes(adev.vram.view(psp.cmd_paddr + RESP, 4)[:]), 0)[0]
        if v != 0xA5A5A5A5: return v
    return None
if sys.argv[1] == 'list':
    for i in range(16):
        print('  %#-9x %s' % (i << 16, NAMES.get(i << 16, 'UNDOCUMENTED')))
    sys.exit(0)
cid = int(sys.argv[1], 0)
print('id=%#x (%s)' % (cid, NAMES.get(cid, 'UNDOCUMENTED')), flush=True)
if not create() or submit() is None:
    print('  CONTROL FAILED -> void'); sys.exit(2)
print('  control before: ring creates and a command completes', flush=True)
h0 = adev.reg(f"{P}_81").read()
adev.reg(f"{P}_64").write(cid)
time.sleep(0.2)
st = adev.reg(f"{P}_64").read()
alive = adev.reg(f"{P}_81").read() != h0
print('  after writing id: _64=%#-12x status=%#-6x hb_advancing=%s' % (st, st & 0xFFFF, alive),
      flush=True)
rec = create()
rs = submit() if rec else None
print('  RESULT %#x (%s): recreate=%s submit=%s -> %s'
      % (cid, NAMES.get(cid, 'UNDOC'), rec, hex(rs) if rs is not None else 'FAIL',
         'PSP STILL WORKS' if rs is not None else 'PSP UNUSABLE AFTERWARDS'), flush=True)