"""Route C decisive probe: does LOAD_IP_FW type 0x55 actually perform the register write?

Prior hardware data (ATTACK-SURFACE.md addendum) recorded type 85 -> 0xffff0006 at
sizes 0x10/0x80/0xFF and was interpreted as "signature verification rejected it".
That reading was WRONG: 0xFFFF0006 is buffer_copy_bounded's own reject, and the old
probe's fill pattern put 0x6f68615a in the address slot -- outside the allow-list
window [0x21500,0x21B78]. So the path had already reached the register-write loop.

Model under test (buffer_copy_bounded @ 0xE0FFD0):
    a1 == 0                    -> 0xFFFF0006
    size > 0x200               -> 0xFFFF300F
    size % 8 != 0              -> 0xFFFF0006
    element addr out of window -> 0xFFFF0006   (loop aborts, no write)
    element addr in window     -> SVC 0x7C write, loop continues, returns 0x0

Stage 1 (ZERO WRITES) validates that model using only rejecting inputs.
Stage 2 performs exactly ONE write, and deliberately uses address 0x21B40, which
buffer_copy_bounded special-cases: on that path the FIRMWARE supplies the value
(qword_E264A8.lo + 0x700000), not us. So the bytes written are exactly what the
firmware writes during a legitimate type-0x55 load -- we control only whether it
happens, not what is written.

Stage 2 is skipped unless --write is passed.
"""
import ctypes, sys, time
from tinygrad import Device
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import lo32, hi32

DO_WRITE = "--write" in sys.argv

d = Device["AMD"]; adev = d.iface.dev_impl; psp = adev.psp
SZ = ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
pref = psp.reg_pref

def alive(): return adev.reg(f"{pref}_81").read() != 0
def stat(): return int.from_bytes(bytes(adev.vram.view(psp.cmd_paddr, SZ)[:])[864:868], "little")

pa = adev.mm.palloc(0x10000, align=0x1000, zero=True)
mc = adev.paddr2mc(pa)

def set_elements(elems):
    """elems = list of (addr, value) 8-byte pairs."""
    buf = bytearray(0x1000)
    for i, (a, v) in enumerate(elems):
        buf[8*i:8*i+4] = (a & 0xFFFFFFFF).to_bytes(4, "little")
        buf[8*i+4:8*i+8] = (v & 0xFFFFFFFF).to_bytes(4, "little")
    adev.vram.view(pa, 0x1000)[:] = bytes(buf)
    adev.gmc.flush_hdp()

def load(fw_type, size):
    adev.vram.view(psp.cmd_paddr, SZ)[:] = b"\x00" * SZ
    c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_IP_FW)
    t = c.cmd.cmd_load_ip_fw
    t.fw_phy_addr_hi, t.fw_phy_addr_lo = hi32(mc), lo32(mc)
    t.fw_size = size; t.fw_type = fw_type
    t0 = time.time()
    try:
        psp._ring_submit(c); s = 0
    except RuntimeError:
        s = stat()
    except TimeoutError:
        return "TIMEOUT", time.time() - t0
    return s, time.time() - t0

def run(label, elems, size, expect):
    if not alive():
        print("  !! sOS DEAD -- aborting"); return None
    set_elements(elems)
    s, dt = load(0x55, size)
    ss = s if isinstance(s, str) else "0x%08x" % s
    ok = "MATCH" if (isinstance(s, int) and s == expect) else ("?" if expect is None else "DIFFERS")
    print("  %-46s -> %-12s %6.3fs alive=%-5s %s" % (label, ss, dt, alive(), ok))
    return s

OUT = 0x6F68615A          # the old probe's pattern: outside the window
IN_SPECIAL = 0x21B40      # in-window AND special-cased -> firmware supplies the value

print("sOS alive at start:", alive())
print("\n=== STAGE 1 - model validation, ZERO writes (every input is rejected) ===")
run("size=0 (zero-length map)",                 [],                 0x000, 0x5)
run("size=0x208 > 0x200",                       [(OUT, 0)]*2,       0x208, 0xFFFF300F)
run("size=0x0C misaligned",                     [(OUT, 0)],         0x00C, 0xFFFF0006)
run("size=8, addr OUT of window (old result)",  [(OUT, 0)],         0x008, 0xFFFF0006)
run("size=8, addr just BELOW window (0x214FC)", [(0x214FC, 0)],     0x008, 0xFFFF0006)
run("size=8, addr just ABOVE window (0x21B7C)", [(0x21B7C, 0)],     0x008, 0xFFFF0006)

print("\n  If all six MATCH, the error model is confirmed and 0x0 can only mean")
print("  the element was accepted and SVC 0x7C executed.")

if not DO_WRITE:
    print("\n=== STAGE 2 skipped (pass --write to run the single write) ===")
else:
    print("\n=== STAGE 2 - ONE write, firmware-chosen value, addr 0x21B40 ===")
    run("size=8, addr 0x21B40 (special-cased)",  [(IN_SPECIAL, 0)], 0x008, None)
    print("  0x0 => the PSP performed the register write, unsigned, host-commanded.")

print("\nsOS alive at end:", alive())