import ctypes, struct, time
from tinygrad import Device
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import lo32, hi32
d=Device["AMD"]; adev=d.iface.dev_impl; psp=adev.psp
SZ=ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
def stat(): return int.from_bytes(bytes(adev.vram.view(psp.cmd_paddr,SZ)[:])[864:868],"little")
def alloc(n):
    pa=adev.mm.palloc(n,align=0x1000,zero=True); return pa,adev.paddr2mc(pa)
P='/Users/user/dev/amd-psp-re/psp_out/'
def load(cmd,blob):
    ap,amc=alloc(0x40000); cp,cmc=alloc(0x8000)
    adev.vram.view(ap,len(blob))[:]=blob; adev.gmc.flush_hdp()
    adev.vram.view(psp.cmd_paddr,SZ)[:]=b"\x00"*SZ
    c=am.struct_psp_gfx_cmd_resp(cmd_id=cmd); t=c.cmd.cmd_load_ta
    t.app_phy_addr_hi,t.app_phy_addr_lo=hi32(amc),lo32(amc); t.app_len=len(blob)
    t.cmd_buf_phy_addr_hi,t.cmd_buf_phy_addr_lo=hi32(cmc),lo32(cmc); t.cmd_buf_len=0x8000
    r=psp._ring_submit(c); return r.resp.session_id,cp
load(am.GFX_CMD_ID_LOAD_ASD, open(P+'asd_full.bin','rb').read())
sid,shm=load(am.GFX_CMD_ID_LOAD_TA, open(P+'hdcp_full.bin','rb').read())
def inv(buf,tc):
    adev.vram.view(shm,0x8000)[:]=bytes(buf); adev.gmc.flush_hdp()
    adev.vram.view(psp.cmd_paddr,SZ)[:]=b"\x00"*SZ
    ic=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD); iv=ic.cmd.cmd_invoke_cmd
    iv.session_id,iv.ta_cmd_id=sid,tc; t0=time.time()
    try: psp._ring_submit(ic); r=0
    except RuntimeError: r=stat()
    except TimeoutError: return "TO",None
    adev.gmc.flush_hdp(); return r, bytes(adev.vram.view(shm,0x8000)[:])
def rd_lenfield():
    b2=bytearray(0x8000); struct.pack_into('<I',b2,0,0x13); r,out=inv(b2,0x13)
    return (struct.unpack_from('<I',out,0x1414)[0] if out else None), (out[0x1418:0x1428].hex() if out else '')
A1P=0x200
for a2p in (2, 0x100, 0x8000, 0xFFFF):
    b2h=(a2p>>8)&0xFF; b3=a2p&0xFF
    b=bytearray(0x8000); struct.pack_into('<I',b,0,0x12); struct.pack_into('<I',b,0xC,0x40)
    b[0x10]=0x91; b[0x11]=0; b[0x12]=b2h; b[0x13]=b3; b[0x14]=0
    b[0x15]=(A1P>>16)&0xFF; b[0x16]=(A1P>>8)&0xFF; b[0x17]=A1P&0xFF
    r1,_=inv(b,0x12)
    ln,hdr=rd_lenfield()
    print("a2'=%#06x -> 0x308 r=%s ; *0xd20438=%s WRITE-BRANCH-FIRED=%s"%(a2p,r1,hex(ln) if ln is not None else None, ln==A1P))
print("sOS alive:",psp.is_sos_alive())