"""Pretty-print exactly what tinygrad's discovery check reads and compares.

_run_discovery() does:
    vram_size  = rreg(mmRCC_CONFIG_MEMSIZE) << 20
    tmr_offset = vram_size - 64KiB
    disc_tbl   = read 10KiB of VRAM at tmr_offset
    assert bhdr.binary_signature == BINARY_SIGNATURE and ihdr.signature == DISCOVERY_TABLE_SIGNATURE
"""
import ctypes, struct
from tinygrad.runtime.support.usb import USB3
from tinygrad.runtime.support.system import USBPCIDevice
from tinygrad.runtime.autogen.am import am

cands = []
for vid, pid in ((0x3801, 0x0001), (0xADD1, 0x0001)):
    cands += list(USB3.list_devices(vid, pid))
gpu = USBPCIDevice("AM", *cands[0])
bar5 = gpu.map_bar(5, fmt='B')
rd32 = lambda off: struct.unpack('<I', bytes(bar5[off:off+4]))[0]

memsize = rd32(0xde3 * 4)
vram_size = memsize << 20
tmr_offset = vram_size - (64 << 10)

def sigstr(v):
    try: return repr(int(v).to_bytes(4, 'little'))
    except Exception: return str(v)

print("mmRCC_CONFIG_MEMSIZE  = %#x" % memsize)
print("vram_size             = %#x (%d MB)" % (vram_size, vram_size >> 20))
print("tmr_offset            = vram_size - 0x10000 = %#x  (%s)"
      % (tmr_offset & 0xFFFFFFFFFFFFFFFF, "NEGATIVE / invalid" if tmr_offset < 0 else "ok"))
print()

bar0 = gpu.map_bar(0)
print("BAR0 (VRAM aperture) size = %#x" % len(bar0))
read_at = tmr_offset if tmr_offset >= 0 else 0
print("reading the discovery table at offset %#x%s" % (read_at, "" if tmr_offset >= 0 else "  (clamped from negative)"))
tbl = bytes(bar0[read_at:read_at + 64])
print("  first 64 bytes: %s" % tbl.hex())
print()

got_bin = struct.unpack_from('<I', tbl, 0)[0]
print("binary_signature        got = %#010x  %s" % (got_bin, sigstr(got_bin)))
print("            am.BINARY_SIGNATURE = %#010x  %s" % (am.BINARY_SIGNATURE, sigstr(am.BINARY_SIGNATURE)))
print("            match: %s" % (got_bin == am.BINARY_SIGNATURE))
print()
print("am.DISCOVERY_TABLE_SIGNATURE = %#010x  %s" % (am.DISCOVERY_TABLE_SIGNATURE, sigstr(am.DISCOVERY_TABLE_SIGNATURE)))
print()
print(">>> The signature check is a SYMPTOM: with MEMSIZE=0 the table offset is")
print(">>> invalid, so it compares against whatever VRAM returns, not a real table.")