#!/usr/bin/env python3
"""Validate the LOAD_TA classification oracle against known ground truth.

The oracle (ccp_classify.py) returns 0x34 for a readable-but-invalid source and
0xA for a source inside the TMR. If that is real rather than coincidence, it
must track the TMR's KNOWN boundary: TMR = PA [0x300000, 0x1600000).

Ground truth comes from the SMU/PSP fields (tmr_paddr=0x300000,
tmr_size=0x1300000), measured previously and independently of this method.

All probes are inside backed VRAM -- no absent addresses (CONSTRAINTS.md).
"""
import sys, time
sys.path.insert(0, "/root/psp/work")
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am
import load_ta

def main():
    adev = Device["AMD"].iface.dev_impl
    hb = adev.reg("regMP0_SMN_C2PMSG_81")
    def alive():
        a = hb.read(); time.sleep(0.25); return hb.read() != a
    if not alive(): raise SystemExit("PSP not healthy before start")
    blob, _ = load_ta.ta_blob("ras")
    cmd_pa = adev.mm.palloc(0x10000, align=0x1000, zero=True, boot=False)
    def probe(pa):
        mc = adev.paddr2mc(pa)
        c = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TA)
        c.cmd.cmd_load_ta.app_phy_addr_lo = mc & 0xFFFFFFFF
        c.cmd.cmd_load_ta.app_phy_addr_hi = mc >> 32
        c.cmd.cmd_load_ta.app_len = len(blob)
        mcc = adev.paddr2mc(cmd_pa)
        c.cmd.cmd_load_ta.cmd_buf_phy_addr_lo = mcc & 0xFFFFFFFF
        c.cmd.cmd_load_ta.cmd_buf_phy_addr_hi = mcc >> 32
        c.cmd.cmd_load_ta.cmd_buf_len = 0x10000
        try:
            r = adev.psp._ring_submit(c); return r.resp.status
        except Exception as e:
            try: return int(str(e).strip().split()[-1])
            except Exception: return -1
    TMR_LO, TMR_HI = 0x300000, 0x1600000
    print("TMR ground truth: [%#x, %#x)" % (TMR_LO, TMR_HI))
    print()
    print("  %-14s %-8s %-10s %s" % ("probe PA", "status", "oracle", "expected"))
    rows = [(TMR_LO - 0x100000, "outside"), (TMR_LO - 0x1000, "outside"),
            (TMR_LO,            "INSIDE"),  (TMR_LO + 0x1000, "INSIDE"),
            (0x900000,          "INSIDE"),  (TMR_HI - 0x10000, "INSIDE"),
            (TMR_HI,            "outside"), (TMR_HI + 0x100000, "outside")]
    agree = 0
    for pa, expect in rows:
        st = probe(pa)
        says = "INSIDE" if st == 0xA else ("outside" if st == 0x34 else "?%#x" % st)
        ok = (says == expect); agree += ok
        print("  %-#14x %-#8x %-10s %-8s %s" % (pa, st, says, expect, "agree" if ok else "MISMATCH"))
    print()
    print("  %d/%d agree with independently-measured TMR bounds" % (agree, len(rows)))
    print("  PSP alive after sweep: %s" % alive())
    return 0 if agree == len(rows) else 1

sys.exit(main())