import struct, sys
from tinygrad import Device
adev = Device["AMD"].iface.dev_impl

# locate the USB controller object (has pcie_mem_read + usb.control_write)
def find_ctrl(root):
    seen=set(); stack=[root]
    while stack:
        o=stack.pop()
        if id(o) in seen: continue
        seen.add(id(o))
        if hasattr(o,'pcie_mem_read') and hasattr(o,'usb'): return o
        for a in ('iface','dev_impl','usb','controller','ctrl','backend','_iface','pci'):
            if hasattr(o,a):
                try: stack.append(getattr(o,a))
                except Exception: pass
    return None
c1=Device["AMD"].iface.ctrl; ctrl=c1 if hasattr(c1,"pcie_mem_read") else adev.doorbell64.usb
print("controller:", type(ctrl).__name__ if ctrl else None)
assert ctrl is not None, "could not find USB controller"

# scratch VRAM page -> PCIe/MC address the TLP engine targets
pa = adev.mm.palloc(0x1000, align=0x1000, zero=True, boot=False)
mc = adev.paddr2mc(pa)
print("scratch paddr=%#x  mc(TLP addr)=%#x" % (pa, mc))

def eb(addr, valA, valB, iters, fmt=0x60, byte_en=0x0F):
    payload = struct.pack('<QII', addr & 0xFFFFFFFFFFFFFFFF, valA & 0xFFFFFFFF, valB & 0xFFFFFFFF)
    payload += bytes([0x48,0x4D,0x52,0x01])   # arming magic 'HMR\x01'
    assert len(payload)==20, len(payload)
    ctrl.usb.control_write(0xEB, value=iters & 0xFFFF, index=(fmt | (byte_en<<8)), data=payload, timeout=5000)

def rd(addr): return struct.unpack('<I', ctrl.pcie_mem_read(addr, 4))[0]

print("before            :", hex(rd(mc)))
eb(mc, 0xCAFEBABE, 0xCAFEBABE, 1)
v1 = rd(mc); print("after A=B=CAFEBABE :", hex(v1), "->", "PASS" if v1==0xCAFEBABE else "FAIL")

eb(mc, 0x11223344, 0xAABBCCDD, 2000)
v2 = rd(mc); print("after A/B x2000    :", hex(v2), "-> last=B=AABBCCDD:", "PASS" if v2==0xAABBCCDD else "FAIL")

# no-magic safety: wrong magic must be a no-op (value stays v2)
bad = struct.pack('<QII', mc, 0xDEAD0000, 0xDEAD0000) + bytes([0,0,0,0])
ctrl.usb.control_write(0xEB, value=1, index=(0x60|0x0F00), data=bad, timeout=5000)
v3 = rd(mc); print("after NO-magic     :", hex(v3), "-> must stay AABBCCDD (gated):", "PASS" if v3==0xAABBCCDD else "FAIL")