#!/usr/bin/env python3
"""Is the LOAD_IP_FW type 49/50 path gated shut, and does loading VCN open it?

FINDING-LOADIPFW-VCN-REGWRITE.md: types 49/50 reach vf_table_lookup -- 64
host-controlled register writes, pre-signature -- but only if
loadfw_engine_step() passes, which needs the VCN enable bit in dword_E257C0.
Their next-step item 2 is exactly "does a host-reachable sequence enable VCN".

This probes the GATE ONLY, with size=0, so NO register write can occur:
  * gate CLOSED -> the engine-step error (-53212 == 0xFFFF3024)
  * gate OPEN   -> falls through to the size/param error (their control: 0x5)

Sequence: probe before any VCN load (control), load VCN ucode, probe again.
Progress written to /tmp/vcngate.txt before each call.
"""
import struct, subprocess, sys, time
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

OUT = "/tmp/vcngate.txt"


def say(s):
    with open(OUT, "a") as f:
        f.write(s + "\n")
    print(s, flush=True)


adev = Device["AMD"].iface.dev_impl
psp = adev.psp
hb = adev.reg("regMP0_SMN_C2PMSG_81")


def alive():
    a = hb.read(); time.sleep(0.3); return hb.read() != a


def probe(fw_type, size, tag):
    """LOAD_IP_FW with a chosen fw_type and fw_size. size=0 -> no writes."""
    cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_IP_FW)
    lo = psp.msg1_addr & 0xFFFFFFFF
    hi = psp.msg1_addr >> 32
    cmd.cmd.cmd_load_ip_fw.fw_phy_addr_lo = lo
    cmd.cmd.cmd_load_ip_fw.fw_phy_addr_hi = hi
    cmd.cmd.cmd_load_ip_fw.fw_size = size
    cmd.cmd.cmd_load_ip_fw.fw_type = fw_type
    try:
        r = psp._ring_submit(cmd)
        st = r.resp.status
    except RuntimeError as e:
        st = int(str(e).rsplit(" ", 1)[-1])
    st &= 0xFFFFFFFF
    say("    %-26s fw_type=%-3d size=%-4d -> status %#010x (%d)" % (tag, fw_type, size, st, st))
    return st


GATE_CLOSED = (-53212) & 0xFFFFFFFF
say("\n######## VCN gate probe ########")
say("expected gate-closed status: %#010x" % GATE_CLOSED)
say("healthy before: %s" % alive())

say("\n=== BEFORE any VCN load (control) ===")
say("  size=0x208 (oversize) rejects INSIDE vf_table_lookup, i.e. AFTER the gate,")
say("  so it separates gate-closed (0xffff3024) from gate-open (0xffff300f).")
before49 = probe(49, 0x208, "VCN0_RAM oversize")
before50 = probe(50, 0x208, "VCN1_RAM oversize")
base85 = probe(85, 0x208, "type 85 oversize (ref)")
probe(49, 0, "VCN0_RAM size0")
probe(85, 0, "type 85 size0 (ref)")
say("  healthy: %s" % alive())

say("\n=== load VCN ucode via LOAD_IP_FW type 0xd ===")
subprocess.run(["zstd", "-d", "-f", "-o", "/tmp/vcn.bin",
                "/lib/firmware/amdgpu/vcn_4_0_0.bin.zst"], check=True,
               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
blob = open("/tmp/vcn.bin", "rb").read()
ua = struct.unpack_from("<I", blob, 0x18)[0]
usz = struct.unpack_from("<I", blob, 0x14)[0]
try:
    psp._load_ip_fw_cmd([am.GFX_FW_TYPE_VCN], memoryview(blob)[ua:ua + usz])
    say("  VCN ucode load: accepted")
except Exception as e:
    say("  VCN ucode load FAILED: %s" % str(e)[:80])
say("  healthy: %s" % alive())

say("\n=== + SMU PowerUpVcn(0) ===")
try:
    adev.smu._send_msg(adev.smu.smu_mod.PPSMC_MSG_PowerUpVcn, 0)
    say("  PowerUpVcn(0): OK")
except Exception as e:
    say("  PowerUpVcn(0) FAILED: %s" % str(e)[:70])
say("  smu alive: %s ; psp alive: %s" % (adev.smu.is_smu_alive(), alive()))
probe(49, 0x208, "VCN0 after powerup")
probe(50, 0x208, "VCN1 after powerup")

say("\n=== AFTER VCN load ===")
after49 = probe(49, 0x208, "VCN0_RAM oversize")
after50 = probe(50, 0x208, "VCN1_RAM oversize")
say("  healthy: %s" % alive())

say("\nVERDICT")
for nm, b, a in (("type 49", before49, after49), ("type 50", before50, after50)):
    if b == a:
        say("  %s unchanged (%#x) -- VCN ucode load does not move this gate" % (nm, b))
    else:
        say("  %s CHANGED %#x -> %#x  <== gate state moved" % (nm, b, a))
say("  reference: Route C type 85 size=0 gave %#x (their control says 0x5)" % base85)