#!/usr/bin/env python3
"""qspi_image.py -- build a whole-flash QSPI image for a Zynq-7000 board.

  python qspi_image.py <out.bin> <offset>:<file> [<offset>:<file> ...]
         [--erase <offset>:<len> ...] [--size <bytes>] [--layout <out.json>]

Every part is placed at its offset; gaps (and --erase ranges, e.g. the U-Boot
environment sector that must stay blank so the default bootcmd is used) are
0xFF, exactly what a fresh `sf erase` leaves.  The image ends at the last byte
of the last part (or --size).  Overlapping parts are refused.  Prints
offset / length / CRC32 of every part and of the image, the numbers to check
against `crc32` in U-Boot, and writes them to --layout as JSON.
Offsets and lengths accept 0x hex.
"""
import json, sys, zlib

def num(s): return int(s, 0)

def main(argv):
    if len(argv) < 3:
        sys.exit(__doc__)
    out, rest = argv[1], argv[2:]
    parts, erase, size, layout = [], [], None, None
    i = 0
    while i < len(rest):
        a = rest[i]
        if a == "--erase":
            o, l = rest[i + 1].split(":"); erase.append((num(o), num(l))); i += 2
        elif a == "--size":
            size = num(rest[i + 1]); i += 2
        elif a == "--layout":
            layout = rest[i + 1]; i += 2
        else:
            o, f = a.split(":", 1); parts.append((num(o), f)); i += 1
    blobs = []
    for o, f in sorted(parts):
        d = open(f, "rb").read()
        blobs.append((o, f, d))
    for (o1, f1, d1), (o2, f2, d2) in zip(blobs, blobs[1:]):
        if o1 + len(d1) > o2:
            sys.exit(f"overlap: {f1} 0x{o1:x}+0x{len(d1):x} runs into {f2} at 0x{o2:x}")
    end = max(o + len(d) for o, _, d in blobs)   # erased tails need no bytes: blank anyway
    if size is not None:
        if size < end:
            sys.exit(f"--size 0x{size:x} smaller than the content end 0x{end:x}")
        end = size
    img = bytearray(b"\xff" * end)
    for o, _, d in blobs:
        img[o:o + len(d)] = d
    for o, l in erase:
        for o2, f2, d2 in blobs:
            if o < o2 + len(d2) and o2 < o + l:
                sys.exit(f"--erase 0x{o:x}:0x{l:x} overlaps {f2}")
    open(out, "wb").write(img)
    rep = {"image": out, "length": len(img), "crc32": "%08x" % zlib.crc32(img), "parts": [],
           "erased": [{"offset": hex(o), "length": hex(l)} for o, l in erase]}
    for o, f, d in blobs:
        rep["parts"].append({"file": f, "offset": hex(o), "length": hex(len(d)),
                             "crc32": "%08x" % zlib.crc32(d)})
        print(f"0x{o:08x}  len 0x{len(d):x}  crc32 {zlib.crc32(d):08x}  {f}")
    for o, l in erase:
        print(f"0x{o:08x}  len 0x{l:x}  erased (0xFF)")
    print(f"image {out}: len 0x{len(img):x}  crc32 {rep['crc32']}")
    if layout:
        json.dump(rep, open(layout, "w"), indent=1)

if __name__ == "__main__":
    main(sys.argv)