#!/usr/bin/env python3
"""Calibrate the SMN->MMIO mapping, then re-run the write diff over the RIGHT window.

Derivation: NV22 SMN 0x1F800 and 0x1E000 divided by 4 give 0x7E00 and 0x7800 --
exactly the two VCN base addresses in this card's own IP discovery table. So
    MMIO dword index = SMN byte address / 4

That puts UVD_AVM_GP_SCRATCH0 (SMN 0x21B04) at dword 0x86C1, OUTSIDE both windows
diffed earlier (0x7800-0x7F00 and 0x7000-0x8200) -- which would explain the
negative.

Step 1 calibrates by writing a scratch register through the host MMIO path and
reading it back (scratch regs exist to be written; harmless).
Step 2 re-runs the PSP write diff over 0x7E00..0x8800.
"""
import struct, time
from tinygrad import Device
import tinygrad.runtime.autogen.am.am as am

OUT = "/tmp/vcncal.txt"
MAGIC = 0xA5A50000


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.25); return hb.read() != a


def d(smn):
    return smn // 4


say("\n######## SMN->MMIO calibration ########")
for smn, nm in ((0x1F844, "DPG_LMA_CTL"), (0x1F848, "DPG_LMA_DATA"),
                (0x1FD58, "UVD_VCPU_CNTL"), (0x1FD00, "VCPU_CACHE_OFFSET0"),
                (0x21B04, "AVM_GP_SCRATCH0"), (0x21B08, "AVM_GP_SCRATCH1")):
    say("   %-20s SMN %#07x -> dword %#06x = %#010x" % (nm, smn, d(smn), adev.rreg(d(smn))))

say("\nscratch write/read-back test (host MMIO path):")
s0 = d(0x21B04)
orig = adev.rreg(s0)
adev.wreg(s0, 0x5A5A5A5A)
rb = adev.rreg(s0)
say("   GP_SCRATCH0 dword %#06x: was %#010x, wrote 0x5A5A5A5A, read %#010x -> %s"
    % (s0, orig, rb, "MAPPING CONFIRMED" if rb == 0x5A5A5A5A else "not a scratch here"))
adev.wreg(s0, orig)

say("\n######## PSP write diff over the corrected window ########")
BASE, END = 0x7E00, 0x8800


def snap():
    o = {}
    for r in range(BASE, END):
        try:
            o[r] = adev.rreg(r)
        except Exception:
            pass
    return o


def submit(payload):
    psp._prep_msg1(memoryview(bytes(payload)))
    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_lo = psp.msg1_addr & 0xFFFFFFFF
    c.cmd.cmd_load_ip_fw.fw_phy_addr_hi = psp.msg1_addr >> 32
    c.cmd.cmd_load_ip_fw.fw_size = len(payload)
    c.cmd.cmd_load_ip_fw.fw_type = 49
    try:
        return psp._ring_submit(c).resp.status & 0xFFFFFFFF
    except RuntimeError as e:
        return int(str(e).rsplit(" ", 1)[-1]) & 0xFFFFFFFF


def ent(i, v):
    return struct.pack("<II", i, v)


before = snap()
say("window %#x..%#x : %d regs, %d non-trivial"
    % (BASE, END, len(before), sum(1 for v in before.values() if v not in (0, 0xFFFFFFFF))))
mid = snap()
say("control (no command): %d regs differ" % len({k for k in mid if before.get(k) != mid[k]}))

# sweep indices across the whole window's worth of dwords
ok = 0
for batch in range(16):
    lo = 0x600 + batch * 63          # 0x600.. covers AVM/ECPU area if index==dword-in-block
    tbl = ent(195, 0) + b"".join(ent(lo + i, MAGIC | ((lo + i) & 0xFFFF)) for i in range(63))
    st = submit(tbl)
    ok += (st == 0)
    if not alive():
        say("  heartbeat stopped at batch %d" % batch); break
after = snap()
dd = {k: (mid.get(k), after[k]) for k in after if mid.get(k) != after[k]}
hits = {k: v for k, (a, v) in dd.items() if (v & 0xFFFF0000) == MAGIC}
say("accepted %d/16 batches ; %d regs changed ; %d hold magic" % (ok, len(dd), len(hits)))
for k, v in list(hits.items())[:12]:
    say("   dword %#06x = %#010x  (SMN %#07x)" % (k, v, k * 4))
say("alive %s" % alive())