#!/usr/bin/env python3
"""Replay the CSU per-field fuzz deterministically and record WHERE it stalls.

Pass 83 saw a stall attributed to cmd 6 / offset 0x3c, which did not reproduce in isolation,
so the cause is cumulative. This replays the identical sequence and prints the running index,
command and field of every invoke, so the stalling step is identified exactly rather than
inferred. Output is flushed per line so the last line before the hang is visible.
"""
import glob, struct, sys, time
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am
adev=Device["AMD"].iface.dev_impl; psp=adev.psp
PAT=0xA5A5A5A5; DECL=0x4000; TOTAL=0x8000
VALUES=[0xFFFFFFFF,0x80000000,0x7FFFFFFF,0x00001000,0xFFFFFFF0]
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(TOTAL)
    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=DECL
    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 inv(sid,shm,cid,off=None,val=0):
    buf=bytearray(TOTAL); struct.pack_into('<I',buf,0,cid)
    if off is not None: struct.pack_into('<I',buf,off,val)
    for o in range(DECL,TOTAL,4): struct.pack_into('<I',buf,o,PAT)
    adev.vram.view(shm,TOTAL)[:]=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
    return st,time.time()-t0
blob=open(glob.glob('/root/psp/contall/ta_42-0-0-6_1c874c3c*')[0],'rb').read()
st,sid,shm=load(blob); print('CSU load %#x sid=%s'%(st,sid), flush=True)
idx=0
for cid in range(0,8):
    base,_=inv(sid,shm,cid); idx+=1
    for off in range(4,0x80,4):
        for v in VALUES:
            s,dt=inv(sid,shm,cid,off,v); idx+=1
            if dt>5:
                print('STALL at invoke #%d: cmd=%d off=%#x val=%#x status=%#x'%(idx,cid,off,v,s), flush=True)
                sys.exit(0)
            if s!=base:
                print('  status change #%d cmd=%d off=%#x val=%#x %#x->%#x'%(idx,cid,off,v,base,s), flush=True)
print('completed %d invokes with no stall'%idx, flush=True)