#!/usr/bin/env python3
"""Re-read the TMR page as RAW BYTES.

Previous run reported "distinct=1008" for the TMR reads, which was an artefact:
the buffer was float32 and NaN != NaN, so set() counts every NaN separately.
0xFFFFFFFF *is* NaN, so that number is consistent with a plain 0xFF fence fill
and proves nothing. Read as uint8 and look at the actual bytes.
"""
import struct, time
import numpy as np
from tinygrad import Device, Tensor
from tinygrad.runtime.support.memory import AddrSpace, PageTableTraverseContext
import tinygrad.runtime.autogen.am.am as am

OUT = "/tmp/tmzraw.txt"


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


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


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


TMR_PA = 0x900000
SZ = 0x1000

src = (Tensor.zeros(SZ, dtype="uint8").contiguous() + 1).contiguous().realize()
va = src.uop.buffer._buf
say("\n######## raw-byte TMR read ########")
say("src VA %#x" % va)

scratch = mm.palloc(SZ, zero=True, boot=False)
adev.vram[scratch:scratch + SZ] = bytes([0x3C]) * SZ
say("scratch PA %#x filled with 0x3C" % scratch)


def remap(pa, tmz):
    mm.unmap_range(va, SZ)
    mm.map_range(va, SZ, [(pa, SZ)], aspace=AddrSpace.PHYS)
    if tmz:
        ctx = PageTableTraverseContext(adev, mm.root_page_table, va, inspect=True)
        for _o, pt, i, cnt, _c in ctx.next(SZ):
            for k in range(cnt):
                pt.entries[i + k] |= am.AMDGPU_PTE_TMZ
    adev.gmc.flush_tlb(ip='GC', vmid=0)


def read(tag):
    try:
        b = (src + 0).realize().numpy().astype(np.uint8)
        u = np.unique(b)
        say("   %-22s first16=%s  unique_bytes=%d %s"
            % (tag, bytes(b[:16]).hex(), len(u), ("all 0x%02x" % u[0]) if len(u) == 1 else "VARIED"))
        return b
    except Exception as e:
        say("   %-22s FAILED %s %s" % (tag, type(e).__name__, str(e)[:60]))
        return None


say("\nhost BAR read of same PA: %s" % bytes(adev.vram[TMR_PA:TMR_PA + 16]).hex())

remap(scratch, False)
c = read("control scratch")
ok = c is not None and len(np.unique(c)) == 1 and c[0] == 0x3C
say("   control valid: %s" % ok)

if ok:
    remap(TMR_PA, False)
    read("TMR no TMZ")
    say("   alive %s" % alive())
    remap(TMR_PA, True)
    read("TMR + TMZ")
    say("   alive %s" % alive())
say("\ndone, alive %s" % alive())