"""Can the HOST invoke ASD directly, choosing the command id?

DISTINCT-WRITE-HUNT.md excluded ASD cmd 0x301 (sub_D132A8 — a clean
content-controlled write of up to 5120 bytes) purely on reachability: no TA
forwarder builds 0x301. Census across all 10 HDCP builds + DTM + RAS confirms
that, so rollback does not open it.

But that argument only covers TA->ASD cross-process RPC. ASD is loaded through
the ordinary TA path (`tee_cmd_dispatch` routes case 1 and case 4 both to
`tee_load_ta`), its header entry is a trampoline
`MOV R0,R2 ; BLX asd_rpc_dispatch ; SVC 0x75`, and asd_rpc_dispatch takes its
opcode from frame[0]. Existing probes call LOAD_ASD but DISCARD the returned
session_id and only ever invoke HDCP.

This probe keeps ASD's session_id and invokes ASD directly, sweeping ta_cmd_id.

Reference points from the dispatcher's own opcode map:
  0x100-0x10D topology/session   0x160-0x161 link      0x180-0x1BF HDCP1.x
  0x1A0-0x1AA HDCP2.2            0x200-0x20E service   0x300-0x309 SRM/revocation
Known-forwarded (so should behave): 0x300 0x302 0x303 0x307 0x308
Target (never forwarded):           0x301
Controls:                           0x999 (no such opcode), 0x200 (service)
"""
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)
P = '/Users/user/dev/amd-psp-re/psp_out/'

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)

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

asd_sid, asd_shm = load(am.GFX_CMD_ID_LOAD_ASD, open(P + 'asd_full.bin', 'rb').read())
print("ASD loaded: session_id=%#x  shm=%#x" % (asd_sid, asd_shm))
if asd_sid == 0:
    print("!! LOAD_ASD returned session_id 0 -- ASD may not be invokable this way")

def invoke_asd(tc, payload=None):
    b = bytearray(0x8000)
    if payload: b[:len(payload)] = payload
    adev.vram.view(asd_shm, 0x8000)[:] = bytes(b); 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 = asd_sid, tc
    t0 = time.time()
    try:
        psp._ring_submit(ic); ps = 0
    except RuntimeError:
        ps = stat()
    except TimeoutError:
        return "TIMEOUT", None, time.time() - t0
    out = bytes(adev.vram.view(asd_shm, 0x40)[:])
    return ps, out, time.time() - t0

print("\n%-28s %-12s %-12s %s" % ("ta_cmd_id", "psp_status", "shm[0:4]", "dt"))
for tc, note in ((0x300, "SRM read (forwarded)"), (0x301, "*** TARGET ***"),
                 (0x302, "SRM query (forwarded)"), (0x303, "SRM check (fwd)"),
                 (0x200, "service/log"), (0x999, "bogus control")):
    ps, out, dt = invoke_asd(tc)
    first = out[:4].hex() if out else "-"
    pss = ps if isinstance(ps, str) else hex(ps)
    print("  %#06x %-20s %-12s %-12s %.3fs" % (tc, note, pss, first, dt))