#!/usr/bin/env python3
"""Does a fence write into the TMR actually land? Read it back directly.

The functional spray was far too sparse (300 dwords over 19 MB) to disturb anything, so it
proves nothing. The TMR is host-allocated VRAM that the PSP protects, so try the direct
question: write a known value via fence_addr and read that offset back from the host.

  * host read shows the value      -> the write LANDED in the trusted region
  * host read blocked/unchanged    -> inconclusive for the write, but tells us about reads

Control: the same procedure against an ordinary scratch page, where the value is known to
land, proves the read-back method works.
"""
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_sz=ctypes.sizeof(am.struct_psp_gfx_rb_frame); step=frame_sz//4
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True,boot=False); return pa,adev.paddr2mc(pa)
tmr_pa=psp.tmr_paddr; tmr_mc=adev.paddr2mc(tmr_pa)
blob=open(glob.glob('/root/psp/contall/ta_42-0-0-6_1c874c3c*')[0],'rb').read()
ap_,amc=alloc(0x80000); cp,cmc=alloc(0x8000); scratch,smc=alloc(0x1000)
n=(len(blob)+3)&~3
adev.vram.view(ap_,n)[:]=blob+b'\x00'*(n-len(blob)); adev.gmc.flush_hdp()
def raw(cmd_bytes, fence_mc, fence_val, wait=0.5):
    adev.vram.view(psp.cmd_paddr, len(cmd_bytes))[:]=cmd_bytes
    adev.vram.view(psp.cmd_paddr+RESP,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=fence_val,
        cmd_buf_addr_lo=adev.paddr2mc(psp.cmd_paddr)&0xFFFFFFFF,
        cmd_buf_addr_hi=adev.paddr2mc(psp.cmd_paddr)>>32,
        fence_addr_lo=fence_mc&0xFFFFFFFF, fence_addr_hi=fence_mc>>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)
    t0=time.time()
    while time.time()-t0<wait:
        adev.gmc.flush_hdp()
        v=struct.unpack_from('<I',bytes(adev.vram.view(psp.cmd_paddr+RESP,4)[:]),0)[0]
        if v!=0xA5A5A5A5: return v
    return None
def mk_load():
    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
    return bytes(memoryview(c).cast('B'))
st=raw(mk_load(), smc, 1)
sid=struct.unpack_from('<I',bytes(adev.vram.view(psp.cmd_paddr+RESP+4,4)[:]),0)[0]
def mk_inv(cid):
    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
    return bytes(memoryview(c).cast('B'))
print('load %s sid=%s; TMR paddr %#x mc %#x'%(hex(st) if st is not None else 'TIMEOUT',sid,tmr_pa,tmr_mc), flush=True)
inv=mk_inv(0)
# can the host read the TMR at all?
pre=[struct.unpack_from('<I',bytes(adev.vram.view(tmr_pa+o,4)[:]),0)[0] for o in (0,0x1000,0x2000,0x100000)]
print('host read of TMR before: %s'%[hex(v) for v in pre], flush=True)
MAGIC=0x5A5AF00D
for tag,off in (('scratch(control)',None),('TMR+0x2000',0x2000),('TMR+0x100000',0x100000)):
    if off is None:
        adev.vram.view(scratch,4,'I')[0]=0; adev.gmc.flush_hdp()
        s=raw(inv, smc, MAGIC); adev.gmc.flush_hdp()
        got=struct.unpack_from('<I',bytes(adev.vram.view(scratch,4)[:]),0)[0]
    else:
        before=struct.unpack_from('<I',bytes(adev.vram.view(tmr_pa+off,4)[:]),0)[0]
        s=raw(inv, tmr_mc+off, MAGIC); adev.gmc.flush_hdp()
        got=struct.unpack_from('<I',bytes(adev.vram.view(tmr_pa+off,4)[:]),0)[0]
        print('    (before=%#x)'%before, flush=True)
    print('  %-18s cmd=%-8s readback=%#-12x %s'%(tag,hex(s) if s is not None else 'TIMEOUT',got,
          'VALUE LANDED' if got==MAGIC else 'not present'), flush=True)