#!/usr/bin/env python3
"""The MP0 FW-debug SMN window decodes from the host. Is it live PSP state?

READ ONLY. A register that decodes is only interesting if it carries
information; a window of constant zeros is decoded and useless. So: dump it,
sample it over time, and sample it across real PSP work (a TA load + invoke).
"""
import struct, time
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

adev = Device["AMD"].iface.dev_impl
psp = adev.psp


def smn_read(addr):
    adev.reg("regBIF_BX0_PCIE_INDEX2").write(addr & 0xFFFFFFFF)
    return adev.reg("regBIF_BX0_PCIE_DATA2").read()


def bswap(x):
    return struct.unpack("<I", struct.pack(">I", x & 0xFFFFFFFF))[0]


BASE = 0x03810000
N = 32


def snap():
    return [smn_read(BASE + i * 4) for i in range(N)]


a = snap()
print("MP0 FW-debug window %#x, %d dwords:" % (BASE, N))
for i in range(0, N, 4):
    print("   +%#04x  %s" % (i * 4, " ".join("%08x" % a[i + j] for j in range(4))))
echo = [i for i in range(N) if a[i] == bswap(BASE + i * 4)]
print("   echoing (not decoded) offsets: %s" % ([hex(i * 4) for i in echo] or "none"))

time.sleep(2.0)
b = snap()
chg = [i for i in range(N) if a[i] != b[i]]
print("\n   changed over 2s idle: %s" % ([hex(i * 4) for i in chg] or "none"))

# now across real PSP work: load the RAS TA and invoke a command
d = open("/tmp/ta.bin", "rb").read()
ua = struct.unpack_from("<I", d, 0x18)[0]
blob = None
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 == 3:
        blob = d[ua + o:ua + o + sz]
before = snap()
app = adev.mm.palloc((len(blob) + 0xFFF) & ~0xFFF, align=0x1000, zero=True, boot=False)
cmdb = adev.mm.palloc(0x10000, align=0x1000, zero=True, boot=False)
adev.vram.view(app, len(blob), "B")[:] = blob
c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA)
mca, mcc = adev.paddr2mc(app), adev.paddr2mc(cmdb)
c.cmd.cmd_load_ta.app_phy_addr_lo = mca & 0xFFFFFFFF
c.cmd.cmd_load_ta.app_phy_addr_hi = mca >> 32
c.cmd.cmd_load_ta.app_len = len(blob)
c.cmd.cmd_load_ta.cmd_buf_phy_addr_lo = mcc & 0xFFFFFFFF
c.cmd.cmd_load_ta.cmd_buf_phy_addr_hi = mcc >> 32
c.cmd.cmd_load_ta.cmd_buf_len = 0x10000
r = psp._ring_submit(c)
sess = r.resp.session_id
shared = adev.vram.view(cmdb, 0x10000, "B")
shared[0:4] = struct.pack("<I", 1)
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 = 1
psp._ring_submit(ic)
after = snap()
uc = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_UNLOAD_TA)
uc.cmd.cmd_unload_ta.session_id = sess
try:
    psp._ring_submit(uc)
except Exception:
    pass

chg2 = [i for i in range(N) if before[i] != after[i]]
print("\n   changed across a real TA load+invoke: %s" % ([hex(i * 4) for i in chg2] or "none"))
for i in chg2:
    print("      +%#04x : %08x -> %08x" % (i * 4, before[i], after[i]))