#!/usr/bin/env python3
"""Confirm the PSP is really RUNNING the substituted (older) SYS_DRV, not silently falling back.

`rollback_boot.py` showed driver init completes and sOS is alive with an older signed SYS_DRV
injected. That is necessary but not sufficient: it shows the bootloader did not refuse, and
not that the older code is what is executing. If the PSP silently kept a resident copy, or if
the load were skipped, the observable state would look the same.

So dump the PSP mailbox registers after boot for each injected build and DIFF them. A version
register that tracks the injected build is direct evidence the substituted component is what
booted. Registers are dumped wholesale rather than assuming which one carries the version.

argv[1]: 'identity' | path to a sos bin
argv[2]: optional label for the output
"""
import ctypes, hashlib, sys
sys.path.insert(0, '/root/psp/hw')
import tinygrad.runtime.support.am.amdev as amdev
import tinygrad.runtime.autogen.am.am as am

def extract_sysdrv(path):
    blob = open(path, 'rb').read()
    hdr = am.struct_psp_firmware_header_v2_0.from_buffer_copy(blob)
    d = am.struct_psp_fw_bin_desc.from_address(ctypes.addressof(hdr.psp_fw_bin))
    if d.fw_type != am.PSP_FW_TYPE_PSP_SYS_DRV:
        raise RuntimeError('first descriptor not SYS_DRV')
    start = d.offset_bytes + hdr.header.ucode_array_offset_bytes
    return blob[start:start + d.size_bytes], d.fw_version

target = sys.argv[1]
label = sys.argv[2] if len(sys.argv) > 2 else target.split('/')[-1]
sub = None; subver = 0
if target != 'identity':
    sub, subver = extract_sysdrv(target)
    assert sub[16:20] == b'$PS1', 'container magic wrong'
orig = amdev.AMFirmware.__init__
def patched(self, adev):
    orig(self, adev)
    running = bytes(self.sos_fw[am.PSP_FW_TYPE_PSP_SYS_DRV])
    newb = running if sub is None else sub
    self.sos_fw[am.PSP_FW_TYPE_PSP_SYS_DRV] = memoryview(bytearray(newb))
    print('INJECT %s ver=%#x len=%d sha=%s'
          % (label, subver, len(newb), hashlib.sha256(newb).hexdigest()[:16]), flush=True)
amdev.AMFirmware.__init__ = patched
from tinygrad import Device
adev = Device["AMD"].iface.dev_impl
P = adev.psp.reg_pref
print('ALIVE %s' % bool(adev.psp.is_sos_alive()), flush=True)
vals = {}
for r in list(range(32, 40)) + list(range(56, 70)) + [81, 92, 115]:
    try: vals[r] = adev.reg(f"{P}_{r}").read()
    except Exception: pass
print('REGS ' + ' '.join('%d=%#x' % (r, v) for r, v in sorted(vals.items())), flush=True)

# TA invocation is serviced BY SYS_DRV, so a working TA proves SYS_DRV is live -- which, from
# a cold start with only the injected bytes supplied, means the injected build is live.
import glob
try:
    psp = adev.psp
    blob = open(glob.glob('/root/psp/contall/ta_42-0-0-6_1c874c3c*')[0], 'rb').read()
    def alloc(n):
        pa = adev.mm.palloc(n, align=0x1000, zero=True, boot=False)
        return pa, adev.paddr2mc(pa)
    ap_, amc = alloc(0x80000); cp, cmc = alloc(0x8000)
    nn = (len(blob) + 3) & ~3
    adev.vram.view(ap_, nn)[:] = blob + b'\x00' * (nn - len(blob)); adev.gmc.flush_hdp()
    c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA); t = c.cmd.cmd_load_ta
    t.app_phy_addr_lo, t.app_phy_addr_hi = amc & 0xFFFFFFFF, amc >> 32; t.app_len = len(blob)
    t.cmd_buf_phy_addr_lo, t.cmd_buf_phy_addr_hi = cmc & 0xFFFFFFFF, cmc >> 32
    t.cmd_buf_len = 0x4000
    r = psp._ring_submit(c)
    st = r.resp.status & 0xFFFFFFFF; sid = r.resp.session_id
    print('TALOAD status=%#x sid=%s' % (st, sid), flush=True)
    if st == 0:
        iv = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_INVOKE_CMD)
        iv.cmd.cmd_invoke_cmd.session_id = sid; iv.cmd.cmd_invoke_cmd.ta_cmd_id = 1
        r2 = psp._ring_submit(iv)
        print('TAINVOKE status=%#x SYSDRV_SERVICING=True' % (r2.resp.status & 0xFFFFFFFF), flush=True)
except Exception as e:
    print('TACHECK ERR %s %s' % (type(e).__name__, str(e)[:120]), flush=True)