#!/usr/bin/env python3
"""Is fence_addr validated against the trusted region? One case per run, with diagnostics.

Established: the PSP writes a host-chosen 4-byte value to a host-chosen address (fence_addr /
fence_value in the ring frame), and `cmd_buf` IS range-checked against the TMR while
fence_addr has never been tested.

argv[1]: scratch | tmr | zero | tmr_end
Reads the mailbox diagnostics afterwards so a hang can be told apart from a fault, the same
way the five known wedges were classified.
"""
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
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
frame_sz=ctypes.sizeof(am.struct_psp_gfx_rb_frame); step=frame_sz//4
def regs(tag):
    d={}
    for r in (33,35,64,65,66,67,81,92):
        try: d[r]=adev.reg(f"{psp.reg_pref}_{r}").read()
        except Exception: d[r]=None
    print('  %-8s alive=%s _81=%#-10x _67=%#-8x _64=%#-10x _33=%#-10x _92=%#x'
          %(tag, bool(d[64] and d[64]&0x80000000), d[81] or 0, d[67] or 0, d[64] or 0,
            d[33] or 0, d[92] or 0), flush=True)
    return d
PAT=0xEEEEEEEE
scratch,smc=alloc(0x1000)
tmr_mc=adev.paddr2mc(psp.tmr_paddr)
blob=open(glob.glob('/root/psp/contall/ta_42-0-0-6_1c874c3c*')[0],'rb').read()
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()
mode=sys.argv[1]
target={'scratch':smc,'tmr':tmr_mc+0x2000,'zero':0,'tmr_end':tmr_mc+0x1300000-0x10}[mode]
print('TMR %#x size 0x1300000; mode=%s target=%#x'%(tmr_mc,mode,target), flush=True)
regs('before')
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=0x4000
# poison the response field so a stale 0 cannot be mistaken for success
adev.vram.view(psp.cmd_paddr, ctypes.sizeof(c))[:]=memoryview(c).cast('B')
adev.vram.view(psp.cmd_paddr+864,4,'I')[0]=0xA5A5A5A5
adev.gmc.flush_hdp()
wmax=psp.ring_size//4
prev=adev.reg(f"{psp.reg_pref}_67").read()%wmax
msg=am.struct_psp_gfx_rb_frame(fence_value=0xDEADBEEF,
    cmd_buf_addr_lo=adev.paddr2mc(psp.cmd_paddr)&0xFFFFFFFF,
    cmd_buf_addr_hi=adev.paddr2mc(psp.cmd_paddr)>>32,
    fence_addr_lo=target&0xFFFFFFFF, fence_addr_hi=target>>32)
adev.vram.view(psp.ring_paddr+prev*4, ctypes.sizeof(msg))[:]=memoryview(msg).cast('B')
adev.reg(f"{psp.reg_pref}_67").write((prev+step)%wmax)
time.sleep(2.0); adev.gmc.flush_hdp()
resp=struct.unpack_from('<I', bytes(adev.vram.view(psp.cmd_paddr+864,4)[:]), 0)[0]
print('  resp field: %#x  (%s)'%(resp,'UNTOUCHED - command never completed' if resp==0xA5A5A5A5
      else 'written by PSP'), flush=True)
regs('after')