#!/usr/bin/env python3
"""Frame offset +8 = `cmd_buf_size`: the one frame field still untested.

In the per-offset frame sweep the +8 row was lost to an inherited dead card and never
retested. It is a LENGTH the host declares for the command buffer -- if the PSP copies that
many bytes into internal memory, an oversized value is an overflow, which is exactly the
shape being hunted.

Per value: does the command complete, does the PSP stay healthy, and is anything written
outside the 4 KB command buffer (guard pages either side, pattern-filled)?

The control is the correct size, and it is re-run after every mutation so a dead PSP cannot
manufacture hits -- the mistake that produced eleven phantom anomalies in the first sweep.
"""
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
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
PAT=0xB7B7B7B7
blob=open(glob.glob('/root/psp/contall/ta_42-0-0-6_1c874c3c*')[0],'rb').read()
ap_,amc=alloc(0x80000); scratch,smc=alloc(0x1000)
# command region: guard | cmd(0x1000) | guard, so an over-read/over-write either side shows
region,rmc=alloc(0x3000)
cmc_cmd=rmc+0x1000
n=(len(blob)+3)&~3
adev.vram.view(ap_,n)[:]=blob+b'\x00'*(n-len(blob)); adev.gmc.flush_hdp()
cbuf,cbmc=alloc(0x8000)
def go(size, tag, wait=1.2):
    adev.vram.view(region,0x3000)[:]=bytes(struct.pack('<I',PAT)*(0x3000//4))
    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
    cb=bytes(memoryview(c).cast('B'))
    adev.vram.view(region+0x1000, len(cb))[:]=cb
    adev.vram.view(region+0x1000+RESP,4,'I')[0]=0xA5A5A5A5
    adev.gmc.flush_hdp()
    f=am.struct_psp_gfx_rb_frame(fence_value=0x1234,
        cmd_buf_addr_lo=cmc_cmd&0xFFFFFFFF, cmd_buf_addr_hi=cmc_cmd>>32,
        fence_addr_lo=smc&0xFFFFFFFF, fence_addr_hi=smc>>32)
    raw=bytearray(memoryview(f).cast('B'))
    if size is not None: struct.pack_into('<I',raw,8,size)
    wmax=psp.ring_size//4; prev=adev.reg(f"{psp.reg_pref}_67").read()%wmax
    adev.vram.view(psp.ring_paddr+prev*4, FRAME)[:]=bytes(raw)
    h0=adev.reg(f"{psp.reg_pref}_81").read()
    adev.reg(f"{psp.reg_pref}_67").write((prev+step)%wmax)
    t0=time.time(); st=None
    while time.time()-t0<wait:
        adev.gmc.flush_hdp()
        v=struct.unpack_from('<I',bytes(adev.vram.view(region+0x1000+RESP,4)[:]),0)[0]
        if v!=0xA5A5A5A5: st=v; break
    adev.gmc.flush_hdp()
    reg=bytes(adev.vram.view(region,0x3000)[:])
    lo=[o for o in range(0,0x1000,4) if struct.unpack_from('<I',reg,o)[0]!=PAT]
    hi=[o for o in range(0x2000,0x3000,4) if struct.unpack_from('<I',reg,o)[0]!=PAT]
    alive=adev.reg(f"{psp.reg_pref}_81").read()!=h0
    return st, lo, hi, alive
st,lo,hi,a=go(None,'control')
print('control (frame as built): resp=%s guard_lo=%d guard_hi=%d alive=%s'
      %(hex(st) if st is not None else 'TIMEOUT',len(lo),len(hi),a), flush=True)
if st is None: print('control failed -> void'); sys.exit(2)
bad=[]
for size in (0,1,4,0x40,0x400,0x1000,0x1001,0x2000,0x10000,0x100000,0x7FFFFFFF,0xFFFFFFFF):
    st,lo,hi,a=go(size,'size')
    cst,clo,chi,ca=go(None,'recheck')
    v=''
    if lo or hi: v='GUARD WRITE lo=%s hi=%s'%([hex(x) for x in lo[:2]],[hex(x) for x in hi[:2]])
    elif st is None and cst is None: v='KILLED THE PSP'
    elif st is None: v='no response, PSP survived'
    print('  size=%#-12x resp=%-9s guards=%d/%d alive=%-5s %s'
          %(size,hex(st) if st is not None else 'TIMEOUT',len(lo),len(hi),a,v), flush=True)
    if v and 'no response' not in v: bad.append((hex(size),v))
    if v=='KILLED THE PSP': break
print('\ncmd_buf_size sweep: %s'%(bad if bad else 'no guard writes, no kills'), flush=True)