"""Validate the IFWI image's OWN internal checksums.

This is independent of how we dumped it: if the image had bitrot, AMD's own
directory checksums would fail. Two independently-obtained dumps agreeing is
suggestive; the image validating its own checksums is much stronger.
"""
import struct, sys, zlib

def fletcher32(data: bytes) -> int:
    s1 = s2 = 0xFFFF
    for i in range(0, len(data) - 1, 2):
        s1 = (s1 + struct.unpack_from('<H', data, i)[0]) % 0xFFFF
        s2 = (s2 + s1) % 0xFFFF
    return (s2 << 16) | s1

def try_all(name, stored, body):
    cands = {
        "fletcher32": fletcher32(body),
        "crc32": zlib.crc32(body) & 0xFFFFFFFF,
        "sum32": sum(struct.unpack_from('<%dI' % (len(body)//4), body, 0)) & 0xFFFFFFFF,
        "sumbytes": sum(body) & 0xFFFFFFFF,
    }
    hit = [k for k, v in cands.items() if v == stored]
    print("    %-22s stored=%#010x  %s" % (name, stored,
          ("MATCH via %s" % hit[0]) if hit else "no match (%s)" %
          ", ".join("%s=%#x" % (k, v) for k, v in cands.items())))
    return bool(hit)

def check(path):
    img = open(path, 'rb').read()
    print("\n=== %s (%d bytes) ===" % (path, len(img)))
    ok = True

    # --- $PSP directories -------------------------------------------------
    for off in (0x2000, 0x3000):
        if img[off:off+4] != b'$PSP':
            print("  %#07x: no $PSP magic (%s)" % (off, img[off:off+4])); ok = False; continue
        stored = struct.unpack_from('<I', img, off+4)[0]
        count  = struct.unpack_from('<I', img, off+8)[0]
        total  = 16 + count*16
        print("  $PSP @ %#07x  entries=%d  total=%#x" % (off, count, total))
        # AMD fletcher32 runs over everything after the checksum field
        if not try_all("dir body (+8..end)", stored, img[off+8:off+total]): ok = False

    # --- ISH headers ------------------------------------------------------
    for off, name in ((0x12000, "ISH A"), (0x13000, "ISH B")):
        w = struct.unpack_from('<8I', img, off)
        stored, gen, poff, psize = w[0], w[1], w[4], w[6]
        print("  %s @ %#07x  gen=%d payload=%#x size=%#x" % (name, off, gen, poff, psize))
        if poff + psize <= len(img):
            if not try_all("payload", stored, img[poff:poff+psize]): ok = False
        else:
            print("    payload out of range"); ok = False
    return ok

paths = sys.argv[1:] or ["psp_out/flash_KNOWNGOOD_20260917.bin", "psp_out/spiflash.bin"]
for p in paths: check(p)