#!/usr/bin/env python

import argparse
import struct


def flip32(data: bytes):
    sl = struct.Struct('<I')
    sb = struct.Struct('>I')
    d = bytearray(len(data))
    for offset in range(0, len(data), 4):
         sb.pack_into(d, offset, sl.unpack_from(data, offset)[0])
    return d


def parse_args():
    parser = argparse.ArgumentParser(description='Convert FPGA bit files to raw bin format suitable for flashing')
    parser.add_argument('-f', '--flip', dest='flip', action='store_true', default=False, help='Flip 32-bit endianess (needed for Zynq)')
    parser.add_argument("bitfiletemplate", help="Input bit file name from which the header is extracted")
    parser.add_argument("binfile", help="Input bin file name from which the bitstream is extracted")
    parser.add_argument("bitfile", help="Output bit file name")
    return parser.parse_args()


def main():
    args = parse_args()
    short = struct.Struct('>H')
    ulong = struct.Struct('>I')

    bitfile = open(args.bitfiletemplate, 'rb')
    headerLength = 0
    l = short.unpack(bitfile.read(2))[0]
    headerLength = headerLength + 2
    if l != 9:
        raise Exception("Missing <0009> header (0x%x), not a bit file" % l)
    bitfile.read(l)
    headerLength = headerLength + l
    l = short.unpack(bitfile.read(2))[0]
    d = bitfile.read(l)
    headerLength = headerLength + l + 2
    if d != b'a':
        raise Exception("Missing <a> header, not a bit file")

    l = short.unpack(bitfile.read(2))[0]
    d = bitfile.read(l)
    headerLength = headerLength + l +2
    print("Design name: %s" % d)

    KEYNAMES = {b'b': "Partname", b'c': "Date", b'd': "Time"}

    while 1:
        k = bitfile.read(1)
        headerLength = headerLength + 1
        if not k:
            bitfile.close()
            raise Exception("unexpected EOF")
        elif k == b'e':
            l = ulong.unpack(bitfile.read(4))[0]
            headerLength = headerLength + 4
            print("Found header end: %d" % headerLength)
            break
        elif k in KEYNAMES:
            l = short.unpack(bitfile.read(2))[0]
            d = bitfile.read(l)
            headerLength = headerLength + l + 2
            print(KEYNAMES[k], d)
        else:
            print("Unexpected key: %s" % k)
            l = short.unpack(bitfile.read(2))[0]
            d = bitfile.read(l)
            headerLength = headerLength + l + 2 

    outputfile = open(args.bitfile, 'wb+')
    bitfile.seek(0)
    outputfile.write(bitfile.read(headerLength))
    bitfile.close()

    binfile = open( args.binfile, 'rb')
    d = binfile.read()
    binfile.close()
    if args.flip:
        print("Flipping data...")
        d = flip32(d)

    outputfile.write(d)
    outputfile.close()


if __name__ == "__main__":
    main()