"""Attempt to boot the card on PSP firmware 0.49.0.51 (a downgrade from the
stock 0.49.0.53 container) and report whether the bootloader accepts it.

This is the first operation in this project that could plausibly have a
persistent effect (if an anti-rollback fuse is enforced), so it is deliberate,
isolated, and reports state before/after.

Requires the PSP to be in a fresh state: run psp_out/psp_recover.py FIRST,
otherwise is_sos_alive() is True and init_hw() loads nothing.

usage: python3 psp_out/test_rollback_51.py [--stock]
   --stock : load the normal .53 container instead (recovery / control run)
"""
import ctypes, sys, hashlib
import tinygrad.runtime.support.am.amdev as amdev
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import mv_address

STOCK = '--stock' in sys.argv
BLOB_51 = '/tmp/lfw/36934f24c008.bin'      # linux-firmware commit with 0.49.0.51

def container_versions(data: bytes):
    blob = memoryview(bytearray(data))
    chdr = am.struct_common_firmware_header.from_address(mv_address(blob))
    H = getattr(am, "struct_psp_firmware_header_v%d_%d" % (chdr.header_version_major, chdr.header_version_minor))
    hdr = H.from_address(mv_address(blob))
    out = {}
    for i in range(hdr.psp_fw_bin_count):
        d = am.struct_psp_fw_bin_desc.from_address(ctypes.addressof(hdr.psp_fw_bin) + i * ctypes.sizeof(am.struct_psp_fw_bin_desc))
        v = d.fw_version
        out[d.fw_type] = "%d.%d.%d.%d" % (v >> 24 & 0xff, v >> 16 & 0xff, v >> 8 & 0xff, v & 0xff)
    return out

if not STOCK:
    data51 = open(BLOB_51, 'rb').read()
    vs = container_versions(data51)
    print("substituting container: SYS_DRV=%s SOS=%s sha=%s" % (
        vs[am.PSP_FW_TYPE_PSP_SYS_DRV], vs[am.PSP_FW_TYPE_PSP_SOS],
        hashlib.sha256(data51).hexdigest()[:12]))
    _orig = amdev.fetch_fw
    def patched(path, name, sha256):
        if name == 'psp_13_0_0_sos.bin':
            print("  -> fetch_fw('%s') INTERCEPTED, returning 0.49.0.51" % name)
            return data51
        return _orig(path, name, sha256)
    amdev.fetch_fw = patched
else:
    print("STOCK run: using the normal pinned container (.53)")

print("initialising device ...")
ok, err = False, None
try:
    from tinygrad import Device
    d = Device["AMD"]
    adev = d.iface.dev_impl
    psp = adev.psp
    ok = psp.is_sos_alive()
    print("RESULT: sOS alive = %s" % ok)
    loaded = container_versions(bytes(adev.fw._raw_sos)) if hasattr(adev.fw, '_raw_sos') else None
    # functional check: a trivial ring command round-trip
    if ok:
        SZ = ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
        adev.vram.view(psp.cmd_paddr, SZ)[:] = b"\x00" * SZ
        c = am.struct_psp_gfx_cmd_resp(cmd_id=0x0f)   # GET_FW_ATTESTATION, read-only
        try:
            psp._ring_submit(c); st = 0
        except RuntimeError:
            st = int.from_bytes(bytes(adev.vram.view(psp.cmd_paddr, SZ)[:])[864:868], "little")
        except TimeoutError:
            st = -1
        print("RESULT: ring cmd 0x0f status = %s" % ("TIMEOUT" if st < 0 else hex(st)))
except Exception as e:
    err = e
    print("RESULT: init FAILED: %r" % (e,))

print("DONE ok=%s err=%s" % (ok, type(err).__name__ if err else None))