#!/usr/bin/env python3
"""Load a signed TA onto the card over the PSP GFX ring, and invoke commands.

tinygrad has no LOAD_TA path, so this builds one. Purpose: drive the RAS TA on
real silicon, to learn whether SYS_DRV context 0xE575B8 -- the capacity field
that is zero in our emulator and blocks RAS opcode 0x000d -- is populated on a
real boot.

    python3 load_ta.py <ta_type> [cmd_id ...]
        ta_type: ras | hdcp | dtm | asd

Health is checked (PSP heartbeat advancing) before and after every submit.
"""
import ctypes, struct, sys, time
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

TY = {"asd": 1, "xgmi": 2, "ras": 3, "hdcp": 4, "dtm": 5}
TA_BIN = "/tmp/ta.bin"

def ta_blob(kind):
    d = open(TA_BIN, "rb").read()
    ua = struct.unpack_from("<I", d, 0x18)[0]
    for i in range(struct.unpack_from("<I", d, 0x20)[0]):
        ft, ver, o, sz = struct.unpack_from("<4I", d, 0x24 + i * 16)
        if ft == TY[kind]:
            return d[ua + o: ua + o + sz], ver
    raise SystemExit("no %s in %s" % (kind, TA_BIN))

def main():
    kind = sys.argv[1] if len(sys.argv) > 1 else "ras"
    cmds = [int(x, 0) for x in sys.argv[2:]] or [1, 2, 3, 4]
    blob, ver = ta_blob(kind)
    print("%s TA: %d bytes, version %d.%d.%d.%d"
          % (kind, len(blob), (ver >> 24) & 0xFF, (ver >> 16) & 0xFF, (ver >> 8) & 0xFF, ver & 0xFF))

    adev = Device["AMD"].iface.dev_impl
    psp = adev.psp
    hb = adev.reg("regMP0_SMN_C2PMSG_81")
    def alive():
        a = hb.read(); time.sleep(0.3); return hb.read() != a
    if not alive():
        raise SystemExit("PSP heartbeat not advancing before we start -- aborting")
    print("PSP healthy before load")

    # TA image buffer + shared command buffer, both in VRAM
    app_pa = adev.mm.palloc(round_up(len(blob)), align=0x1000, zero=True, boot=False)
    cmd_pa = adev.mm.palloc(0x10000, align=0x1000, zero=True, boot=False)
    adev.vram.view(app_pa, len(blob), "B")[:] = blob
    print("TA image  @ PA %#x (MC %#x)" % (app_pa, adev.paddr2mc(app_pa)))
    print("cmd buf   @ PA %#x (MC %#x)" % (cmd_pa, adev.paddr2mc(cmd_pa)))

    c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA)
    mc_app, mc_cmd = adev.paddr2mc(app_pa), adev.paddr2mc(cmd_pa)
    c.cmd.cmd_load_ta.app_phy_addr_lo = mc_app & 0xFFFFFFFF
    c.cmd.cmd_load_ta.app_phy_addr_hi = mc_app >> 32
    c.cmd.cmd_load_ta.app_len = len(blob)
    c.cmd.cmd_load_ta.cmd_buf_phy_addr_lo = mc_cmd & 0xFFFFFFFF
    c.cmd.cmd_load_ta.cmd_buf_phy_addr_hi = mc_cmd >> 32
    c.cmd.cmd_load_ta.cmd_buf_len = 0x10000
    try:
        r = psp._ring_submit(c)
        status, sess = r.resp.status, r.resp.session_id
    except RuntimeError as e:
        print("LOAD_TA failed: %s" % e)
        print("PSP alive after failure: %s" % alive())
        return 1
    print("LOAD_TA -> status=%#x session_id=%#x" % (status, sess))
    if not alive():
        print("!! PSP heartbeat stopped after LOAD_TA"); return 2

    shared = adev.vram.view(cmd_pa, 0x10000, "B")
    # The TA verdict is NOT resp.status -- that is ring delivery, and it returns
    # 0 for a bogus command id too (verified as a control). For RAS the shared
    # buffer is struct ta_ras_shared_memory:
    #     +0x00 cmd_id  +0x04 resp_id  +0x08 ras_status  +0x0c if_version
    #     +0x10 ret_status
    # The command id goes IN THE BUFFER, the same shape as the HDCP ABI.
    print("  %-14s %-8s %-12s %-12s %s" % ("cmd", "ring", "ras_status", "resp_id", "ret"))
    for cid in cmds:
        shared[:0x100] = bytes(0x100)
        shared[0:4] = struct.pack("<I", cid)
        ic = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD)
        ic.cmd.cmd_invoke_cmd.session_id = sess
        ic.cmd.cmd_invoke_cmd.ta_cmd_id = cid
        try:
            rr = psp._ring_submit(ic)
            st = "%#x" % rr.resp.status
        except RuntimeError as e:
            st = str(e).rsplit(" ", 1)[-1]
        hdr = bytes(shared[0:20])
        _cid, resp_id, ras_status, _ifv, retw = struct.unpack("<5I", hdr)
        ret = retw & 0xFF
        print("  cmd %#08x   %-8s %-#12x %-#12x %d" % (cid, st, ras_status, resp_id, ret))
        if not alive():
            print("  !! PSP heartbeat stopped after cmd %#x" % cid); return 2
    # Unload the session. LOAD_TA allocates a new session id every call and
    # they were observed climbing 0,1,2,... across runs, so leaving them behind
    # leaks a finite kernel resource.
    uc = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA)
    uc.cmd.cmd_unload_ta.session_id = sess
    try:
        ur = psp._ring_submit(uc)
        print("UNLOAD_TA session %#x -> status=%#x" % (sess, ur.resp.status))
    except RuntimeError as e:
        print("UNLOAD_TA session %#x -> %s" % (sess, str(e).rsplit(" ", 1)[-1]))
    print("PSP healthy after all invokes: %s" % alive())
    return 0

def round_up(n, a=0x1000): return (n + a - 1) & ~(a - 1)

if __name__ == "__main__":
    raise SystemExit(main())