#!/usr/bin/env python3
"""Pack a raw RP2350 flash image (objcopy -O binary output) into a UF2.

Family rp2350-arm-s (0xe48bff59), flash base 0x10000000. Usage:
    python3 uf2.py pm-kit.bin pm-kit.uf2
"""
import struct
import sys

BASE = 0x10000000
FAMILY = 0xE48BFF59  # rp2350-arm-s

src = sys.argv[1] if len(sys.argv) > 1 else "pm-kit.bin"
out = sys.argv[2] if len(sys.argv) > 2 else "pm-kit.uf2"

data = open(src, "rb").read()
chunks = [data[i:i + 256] for i in range(0, len(data), 256)] or [b""]
n = len(chunks)
with open(out, "wb") as f:
    for i, c in enumerate(chunks):
        c = c.ljust(256, b"\x00")
        blk = struct.pack("<IIIIIIII", 0x0A324655, 0x9E5D5157, 0x00002000,
                          BASE + i * 256, 256, i, n, FAMILY)
        blk += c + b"\x00" * (476 - 256) + struct.pack("<I", 0x0AB16F30)
        assert len(blk) == 512
        f.write(blk)
print(f"{out}: {n} blocks, {len(data)} bytes payload")