#!/usr/bin/env python3
"""Zero-write every stack slot in GCP's live frames; record survive/fault + fault signature.

Purpose: find control-flow divergence or error paths we have NOT theorised. For each target the
4x STRB window is aimed at P = T, zeroing the whole word T..T+3, with a FRESH TA load per slot so
tests are independent.

Measured frame map (22.0.4.50, hardware-confirmed in Addendum 6):
   stack [0x415000,0x419000)  SP0 = 0x419000   SP at the store = 0x418fa0
   gcp_dispatch  R4 0x418ff0  R5 0x418ff4  R6 0x418ff8  LR 0x418ffc (=0x405cdc, ARM)
   dispatcher    R3 0x418fc8 .. R11 0x418fe8  LR 0x418fec (=0x404053)
   handler FE09  R3 0x418fa0 .. R11 0x418fc0  LR 0x418fc4 (=0x402a3b)

A fault wedges the sOS ring, so the run STOPS at the first fault and prints RESUME=<next index>;
the outer driver power-cycles and restarts there. Non-faulting slots cost nothing.

argv[1] = start index. State is appended to /tmp/slotsweep.log on the box.
"""
import struct, sys, glob, os
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
PARAMS_VA=0x04013000; BASE=PARAMS_VA+0x418
TA='ta_22-0-4-50_81b082b4'; BUFLEN=0x8000; PAT=0xA5
STORE_LO,STORE_HI=0x402df0,0x402e10
FAULT=[getattr(adev,'mmMP0_SMN_C2PMSG_%d'%n) for n in range(82,89)]
def frame(): return tuple(r.read()&0xFFFFFFFF for r in FAULT)
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
app_pa,app_mc=alloc(0x80000); cp,cmc=alloc(BUFLEN)
blob=open(glob.glob('/root/psp/contall/%s*'%TA)[0],'rb').read()
def load():
    n=(len(blob)+3)&~3
    adev.vram.view(app_pa,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=app_mc&0xFFFFFFFF,app_mc>>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=BUFLEN
    try:
        r=psp._ring_submit(c); return r.resp.status&0xFFFFFFFF, r.resp.session_id
    except Exception: return -1,None
def unload(sid):
    c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA); c.cmd.cmd_unload_ta.session_id=sid
    try: return psp._ring_submit(c).resp.status&0xFFFFFFFF
    except Exception: return -1
def fe09(sid,disp):
    buf=bytearray(bytes([PAT])*BUFLEN)
    struct.pack_into('<I',buf,0,0xFE09); struct.pack_into('<I',buf,0x10,0)
    for o in (0x414,0x818): struct.pack_into('<I',buf,o,0xFFFFFFFB)
    struct.pack_into('<I',buf,0x820,1); struct.pack_into('<I',buf,0x828,disp&0xFFFFFFFF)
    adev.vram.view(cp,BUFLEN)[:]=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,0xFE09
    try: st=psp._ring_submit(c).resp.status&0xFFFFFFFF
    except Exception: st=-1
    adev.gmc.flush_hdp()
    return st, struct.unpack_from('<I',bytes(adev.vram.view(cp,16)[:]),8)[0]

OWN={}
for i in range(10):
    OWN[0x418fa0+4*i]=('handler', ['R3','R4','R5','R6','R7','R8','R9','R10','R11','LR'][i])
for i in range(10):
    OWN[0x418fc8+4*i]=('dispatcher', ['R3','R4','R5','R6','R7','R8','R9','R10','R11','LR'][i])
for i,r in enumerate(['R4','R5','R6','LR']):
    OWN[0x418ff0+4*i]=('gcp_dispatch', r)
TARGETS=[]
for a in sorted(OWN): TARGETS.append((a,'%s %s'%OWN[a]))
for a,lbl in ((0x418f9c,'dead: just below SP'),(0x418f00,'dead: SP-0xa0'),
              (0x418800,'dead: mid stack'),(0x415004,'dead: stack floor+4'),
              (0x415000,'stack floor'),(0x414ffc,'GUARD below stack'),
              (0x419000,'GUARD at stack top'),(0x418ffc+4,'=0x419000 dup')):
    TARGETS.append((a,lbl))

start=int(sys.argv[1]) if len(sys.argv)>1 else 0
f0=frame()
log=open('/tmp/slotsweep.log','a')
print('baseline frame %s  start=%d/%d'%([hex(x) for x in f0],start,len(TARGETS)),flush=True)
for i in range(start,len(TARGETS)):
    T,lbl=TARGETS[i]
    st,sid=load()
    if sid is None:
        print('IDX %d: load failed (%s) -- ring likely wedged'%(i,st),flush=True)
        print('RESUME=%d'%i,flush=True); break
    D=(T-BASE)&0xFFFFFFFF
    st2,ta=fe09(sid,D)
    fr=frame()
    faulted=(st2==-1) or (fr!=f0)
    if not faulted:
        line='IDX %-2d %#08x %-22s D=%#010x SURVIVED ta_status=%#x'%(i,T,lbl,D,ta)
        print(line,flush=True); log.write(line+'\n'); log.flush()
        unload(sid)
        continue
    pc,w3,addr=fr[2],fr[3],fr[4]
    kind='FAULT@STORE' if STORE_LO<=pc<=STORE_HI else ('RET-TO-ZERO' if (pc==0 or addr==0) else 'OTHER')
    line=('IDX %-2d %#08x %-22s D=%#010x FAULT %s pc=%#x w3=%#x addr=%#x  full=%s'
          %(i,T,lbl,D,kind,pc,w3,addr,[hex(x) for x in fr]))
    print(line,flush=True); log.write(line+'\n'); log.flush()
    print('RESUME=%d'%(i+1),flush=True)
    break
else:
    print('SWEEP-COMPLETE',flush=True)