"""GET_FW_ATTESTATION (0x0F cmd_id) output inspection + resp descriptor handling.

Host-reachable, returns real data, never examined. Two questions:
 1) does the response contain anything beyond the documented measurement
    (PSP-range pointers / uninitialised bytes)?
 2) is the host-supplied resp_offset / resp_buf_size descriptor bounds-checked,
    or will the PSP write wherever we point it?
"""
import ctypes,struct
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)
RANGES=[(0x00020000,0x00070000,'sOS'),(0x00400000,0x00412000,'HDCP'),
        (0x00D00000,0x00D50000,'ASD'),(0x00E10000,0x00E80000,'SYSDRV'),(0x0020A000,0x0020C000,'TAstage')]
def classify(v):
    for lo,hi,nm in RANGES:
        if lo<=v<hi: return nm
    return None
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)
rb_pa,RB=alloc(0x10000)
MARK=0x5A
def go(resp_off=0, resp_size=0x1000, use_rb=True):
    adev.vram.view(rb_pa,0x10000)[:]=bytes([MARK])*0x10000; adev.gmc.flush_hdp()
    b=bytearray(SZ); struct.pack_into('<I',b,8,am.GFX_CMD_ID_GET_FW_ATTESTATION)
    if use_rb:
        struct.pack_into('<I',b,12,lo32(RB)); struct.pack_into('<I',b,16,hi32(RB))
        struct.pack_into('<I',b,20,resp_off); struct.pack_into('<I',b,24,resp_size)
    adev.vram.view(psp.cmd_paddr,SZ)[:]=bytes(b)
    c=am.struct_psp_gfx_cmd_resp.from_buffer_copy(bytes(b))
    try: psp._ring_submit(c); st=0
    except RuntimeError: st=stat()
    except TimeoutError: return "TIMEOUT",None,None
    out=bytes(adev.vram.view(psp.cmd_paddr,SZ)[:])
    rb=bytes(adev.vram.view(rb_pa,0x10000)[:])
    return st,out,rb
st,out,rb=go()
print("GET_FW_ATTESTATION status=%s"%(hex(st) if isinstance(st,int) else st))
if out:
    print("  cmd-buffer response bytes 864..920:", out[864:920].hex())
    nz=[i for i in range(0x40,864) if out[i]]
    print("  non-zero request-area bytes: %d"%len(nz))
    ptr=[(hex(o),hex(struct.unpack_from('<I',out,o)[0]),classify(struct.unpack_from('<I',out,o)[0]))
         for o in range(0,SZ-3,4) if classify(struct.unpack_from('<I',out,o)[0])]
    print("  PSP-range dwords in cmd buffer: %s"%(ptr if ptr else 'NONE'))
if rb:
    wrote=[i for i in range(len(rb)) if rb[i]!=MARK]
    print("  resp buffer: PSP wrote %d bytes%s"%(len(wrote), (" at %#x..%#x"%(wrote[0],wrote[-1])) if wrote else ""))
    if wrote:
        print("   first 64B:", rb[wrote[0]:wrote[0]+64].hex())
        p=[(hex(o),hex(struct.unpack_from('<I',rb,o)[0]),classify(struct.unpack_from('<I',rb,o)[0]))
           for o in range(0,0x2000-3,4) if classify(struct.unpack_from('<I',rb,o)[0])]
        print("   PSP-range dwords in resp buffer: %s"%(p[:10] if p else 'NONE'))
print("\n--- resp_offset bounds check: declare size 0x40 but offset far beyond ---")
for off,size in ((0,0x1000),(0x40,0x40),(0x8000,0x40),(0xfff0,0x40),(0xffffffff,0x40),(0,0)):
    st,out,rb=go(off,size)
    wrote=[i for i in range(len(rb)) if rb[i]!=MARK] if rb else []
    print("  resp_offset=%#010x size=%#06x -> status=%-6s wrote %d bytes%s"%(
        off,size,hex(st) if isinstance(st,int) else st,len(wrote),
        (" at %#x..%#x"%(wrote[0],wrote[-1])) if wrote else ""))
print("sos_alive:",psp.is_sos_alive())