import struct
from tinygrad import Device
adev = Device["AMD"].iface.dev_impl
ctrl = adev.doorbell64.usb
CFG=0x0009
def rd(mc): return struct.unpack('<I', ctrl.pcie_mem_read(mc,4))[0]
def page(): 
    pa=adev.mm.palloc(0x1000, align=0x1000, zero=True, boot=False); return adev.paddr2mc(pa)
def trig(): ctrl.usb.control_write(0xEB, data=bytes([0x48,0x4D,0x52,0x01]), timeout=5000)
def cfg(db_addr, db_val, tgt, A, B, iters):
    c=struct.pack('<QIQIIHBB', db_addr, db_val, tgt, A, B, iters, 0x60, 0x0F)
    ctrl.write(CFG, c); assert bytes(ctrl.read(CFG,32))==c

d1, t1 = page(), page()
print("Test A (doorbell only, iters=0): db=%#x tgt=%#x"%(d1,t1))
cfg(d1, 0xD00DB011, t1, 0xAAAAAAAA, 0xBBBBBBBB, 0); trig()
print("  db=%#x (want D00DB011 %s)  tgt=%#x (want 0 %s)"%(rd(d1),"OK" if rd(d1)==0xD00DB011 else "BAD", rd(t1), "OK" if rd(t1)==0 else "BAD"))

d2, t2 = page(), page()
print("Test B (hammer only, doorbell=0): db=%#x tgt=%#x"%(d2,t2))
cfg(0, 0, t2, 0xAAAAAAAA, 0xBBBBBBBB, 50); trig()
print("  db=%#x (want 0 %s)  tgt=%#x (want BBBBBBBB %s)"%(rd(d2),"OK" if rd(d2)==0 else "BAD", rd(t2), "OK" if rd(t2)==0xBBBBBBBB else "BAD"))

d3, t3 = page(), page()
print("Test C (both): db=%#x tgt=%#x"%(d3,t3))
cfg(d3, 0xD00DB011, t3, 0xAAAAAAAA, 0xBBBBBBBB, 50); trig()
print("  db=%#x (want D00DB011 %s)  tgt=%#x (want BBBBBBBB %s)"%(rd(d3),"OK" if rd(d3)==0xD00DB011 else "BAD", rd(t3), "OK" if rd(t3)==0xBBBBBBBB else "BAD"))
import os,sys; sys.stdout.flush(); os._exit(0)