#!/usr/bin/env python3
"""The ring frame's `vmid` and `frame_type` bytes (offset 32/33) -- host-controlled.

The dword at +32 killed the PSP. In `psp_gfx_rb_frame` that offset is not reserved: it is
  +32 vmid (u8), +33 frame_type (u8), +34..35 reserved1
so the host chooses the VMID the PSP uses to translate the addresses in the frame.

Two questions:
  1. which byte kills it -- vmid or frame_type?
  2. does the PSP HONOUR a chosen vmid? If a non-zero vmid still writes the fence to our
     page, vmid is ignored. If the fence stops landing, the PSP is translating through a
     host-chosen address space -- which in a virtualized part is another VM's.

argv: <vmid> <frame_type>   (-1 = leave at default)
Each run is a fresh boot; the caller recovers if the card wedges.
"""
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)
MAGIC=0x7E571234
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()
vmid=int(sys.argv[1]); ftype=int(sys.argv[2]) if len(sys.argv)>2 else -1
def go(cmd_bytes, tag, wait=1.0):
    adev.vram.view(scratch,4,'I')[0]=0; adev.gmc.flush_hdp()
    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()
    f=am.struct_psp_gfx_rb_frame(fence_value=MAGIC,
        cmd_buf_addr_lo=adev.paddr2mc(psp.cmd_paddr)&0xFFFFFFFF,
        cmd_buf_addr_hi=adev.paddr2mc(psp.cmd_paddr)>>32,
        fence_addr_lo=smc&0xFFFFFFFF, fence_addr_hi=smc>>32)
    raw=bytearray(memoryview(f).cast('B'))
    if vmid>=0: raw[32]=vmid & 0xFF
    if ftype>=0: raw[33]=ftype & 0xFF
    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(psp.cmd_paddr+RESP,4)[:]),0)[0]
        if v!=0xA5A5A5A5: st=v; break
    adev.gmc.flush_hdp()
    fence=struct.unpack_from('<I',bytes(adev.vram.view(scratch,4)[:]),0)[0]
    alive=adev.reg(f"{psp.reg_pref}_81").read()!=h0
    print('  %-22s vmid=%-4s ftype=%-4s resp=%-9s fence_landed=%-5s hb_advancing=%s'
          %(tag, vmid if vmid>=0 else 'dflt', ftype if ftype>=0 else 'dflt',
            hex(st) if st is not None else 'TIMEOUT', fence==MAGIC, alive), flush=True)
    return st, fence==MAGIC, alive
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
go(bytes(memoryview(c).cast('B')), 'LOAD_TA')