"""Is the IP discovery table actually present in VRAM, despite MEMSIZE=0?

If memory works and only RCC_CONFIG_MEMSIZE is unreported, the table still sits
at (vram_size - 64KiB) and we can just tell tinygrad the right size.

READ-ONLY.
"""
import struct
from tinygrad.runtime.support.usb import USB3
from tinygrad.runtime.support.system import USBPCIDevice
from tinygrad.runtime.autogen.am import am

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

print("mmRCC_CONFIG_MEMSIZE = %#x" % rd32(0xde3*4))
print("BAR0 aperture size   = %#x (%d GB)" % (len(bar0), len(bar0) >> 30))
print("looking for BINARY_SIGNATURE = %#010x (%r)"
      % (am.BINARY_SIGNATURE, am.BINARY_SIGNATURE.to_bytes(4, 'little')))
print()

# candidate VRAM sizes for a 7900 XTX and friends
cands_mb = [24576, 20480, 16384, 12288, 8192, 4096, 2048]
found = []
for mb in cands_mb:
    size = mb << 20
    off = size - (64 << 10)
    if off + 16 > len(bar0):
        print("  %6d MB -> offset %#012x  (beyond BAR aperture, skipped)" % (mb, off)); continue
    try:
        b = bytes(bar0[off:off+16])
    except Exception as e:
        print("  %6d MB -> read failed: %s" % (mb, e)); continue
    sig = struct.unpack_from('<I', b, 0)[0]
    hit = (sig == am.BINARY_SIGNATURE)
    print("  %6d MB -> offset %#012x  first16=%s  %s" % (mb, off, b.hex(), "*** MATCH ***" if hit else ""))
    if hit: found.append(mb)

print()
# also sweep the low part of the aperture for the signature, in case the table
# landed somewhere unexpected
print("sweeping first 64 MB of the aperture for the signature (4KiB steps)...")
hits = []
for off in range(0, 64 << 20, 4096):
    try:
        if struct.unpack_from('<I', bytes(bar0[off:off+4]), 0)[0] == am.BINARY_SIGNATURE:
            hits.append(off)
            if len(hits) > 8: break
    except Exception:
        break
print("  signature hits: %s" % ([hex(h) for h in hits] if hits else "none"))

print()
if found:
    print(">>> discovery table FOUND assuming vram_size = %d MB" % found[0])
    print(">>> memory is usable; MEMSIZE reporting is the only broken part.")
else:
    print(">>> no discovery table anywhere -> VRAM is genuinely not up.")
    print(">>> overriding MEMSIZE would not help.")