#!/usr/bin/env python3
"""Repack directory (headers.json + sect*.bin) -> .bin. Verifies exact-fit chain."""
import argparse, json, glob, os, struct, sys
sys.path.insert(0, os.path.dirname(__file__))
from lofi_image import build_image

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('indir')
    ap.add_argument('output')
    ap.add_argument('--checksum', default=None,
                    help="'auto' (default: compute valid checksum), hex u32, or 'keep' (headers.json value)")
    a = ap.parse_args()
    meta = json.load(open(os.path.join(a.indir, 'headers.json')))
    # cmtd raw hex -> rebuild path needs raw_0x0C_0x30; headers.json stores it
    files = sorted(glob.glob(os.path.join(a.indir, 'sect*_addr*.bin')))
    assert files, 'no sect files found'
    payloads = []
    for f in files:
        base = os.path.basename(f)
        addr = int(base.split('addr')[1].split('.')[0], 16)
        payloads.append((addr, open(f, 'rb').read()))
    if a.checksum is None or a.checksum == 'auto':
        ck = None
    elif a.checksum == 'keep':
        ck = 'keep'
    else:
        ck = int(a.checksum, 16)
    out = build_image(meta, payloads, checksum=ck)
    open(a.output, 'wb').write(out)
    print(f"wrote {a.output} ({len(out)} B) checksum={struct.unpack('<I', out[8:12])[0]:08X}")

if __name__ == '__main__':
    main()