#!/usr/bin/env python3
"""Isolate AUC command 6: does it reproducibly hang the PSP?

The twelve-TA silicon fuzz died at AUC cmd 6 -- 10s stall, then every later ring submit
returned an incrementing error and no further TA would load, i.e. the PSP stopped servicing
the ring. The card needed an ATX cycle. This re-runs it alone, with cmds 4 and 5 first as
in-run controls, to separate "cmd 6 specifically" from "AUC after N commands" or harness drift.
"""
import glob, os, struct, sys, time
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am
adev = Device["AMD"].iface.dev_impl; psp = adev.psp
def alloc(n):
    pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False); return pa, adev.paddr2mc(pa)
def load(blob):
    ap_, amc = alloc(0x80000); cp, cmc = alloc(0x8000)
    n=(len(blob)+3)&~3
    adev.vram.view(ap_,n)[:] = blob + b'\x00'*(n-len(blob)); adev.gmc.flush_hdp()
    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 = cmc&0xFFFFFFFF, cmc>>32; t.cmd_buf_len=0x8000
    try:
        r=psp._ring_submit(c); return r.resp.status&0xFFFFFFFF, r.resp.session_id, cp
    except Exception as e:
        try: return int(str(e).strip().split()[-1])&0xFFFFFFFF, None, cp
        except Exception: return -1, None, cp
def invoke(sid, shm, cid):
    buf=bytearray(0x2000); struct.pack_into('<I',buf,0,cid)
    adev.vram.view(shm,len(buf))[:]=bytes(buf); adev.gmc.flush_hdp()
    c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD)
    iv=c.cmd.cmd_invoke_cmd; iv.session_id, iv.ta_cmd_id = sid, cid
    t0=time.time()
    try:
        r=psp._ring_submit(c); st=r.resp.status&0xFFFFFFFF
    except Exception as e:
        try: st=int(str(e).strip().split()[-1])&0xFFFFFFFF
        except Exception: st=-1
    dt=time.time()-t0
    try:
        adev.gmc.flush_hdp(); out=bytes(adev.vram.view(shm,16)[:])
    except Exception: out=b''
    return st, dt, out
cands=glob.glob('/root/psp/contall/ta_16-0-0-9_9e9044cf*')
blob=open(cands[0],'rb').read()
st,sid,shm=load(blob)
print('AUC load: %#x sid=%s'%(st,sid))
if st!=0: sys.exit(2)
for cid in (4,5,6):
    s,dt,out=invoke(sid,shm,cid)
    print('  cmd %-2d -> status %#-12x %5.1fs out=%s'%(cid,s,dt,out[:8].hex()))
    if dt>5: print('     ^^ STALLED'); break
# is the ring still alive?
s,dt,out=invoke(sid,shm,1)
print('  re-probe cmd 1 -> status %#x (%.1fs)  ring alive: %s'%(s,dt,'no' if dt>5 else 'yes'))