#!/usr/bin/env python3
"""extract_fft.py — extract FFT-related code regions from the rt snap into raw bins."""
import struct, mmap, os

SNAP = '/tmp/snap_rt.bin'
OUT = '/tmp/fft/'

# (name, addr, size)
TARGETS = [
    ('cplx_mul_8440', 0x180008440, 0x80),
    ('cplx_mul_kernel_c440', 0x18000c440, 0x100),
    ('stage_bfc0', 0x18000bfc0, 0x600),
    ('stage_c5e0', 0x18000c5e0, 0x600),
    ('twiddle_loader_39b00', 0x180039b00, 0x400),
    ('plan_gen_2f980', 0x18002f980, 0x1200),
    ('dispatcher_535a70', 0x180535a70, 0x100),
    ('scalar_140a10', 0x180140a10, 0x200),
    ('vector_140a70', 0x180140a70, 0x200),
    ('fft_kernel_a5a0', 0x18001a5a0, 0x100),
]


def main():
    os.makedirs(OUT, exist_ok=True)
    fd = os.open(SNAP, os.O_RDONLY)
    sz = os.fstat(fd).st_size
    mm = mmap.mmap(fd, 0, access=mmap.ACCESS_READ)
    # parse region index once
    regs = []
    i = 0
    while i + 16 <= sz:
        lo, n = struct.unpack_from('<QQ', mm, i)
        regs.append((lo, i + 16, n))  # (addr, data_off, size)
        i += 16 + n
    regs.sort()

    def readabs(addr, n):
        for lo, off, rn in regs:
            if lo <= addr < lo + rn and addr - lo + n <= rn:
                return mm[off + (addr - lo): off + (addr - lo) + n]
        return None

    for name, addr, n in TARGETS:
        b = readabs(addr, n)
        if b:
            with open(os.path.join(OUT, name + '.bin'), 'wb') as f:
                f.write(b)
            print('%-28s 0x%x %d bytes OK' % (name, addr, len(b)))
        else:
            print('%-28s 0x%x MISSING' % (name, addr))
    mm.close()
    os.close(fd)


if __name__ == '__main__':
    main()