#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
make_uf2.py — 最小 UF2 封装器(RP2040)
==================================================
背景: pico-sdk 2.x 生成 UF2 依赖 picotool(FetchContent 拉 github), 本机
网络对 github 不通时走本脚本替代: elf 先经 arm-none-eabi-objcopy 转 bin,
再按 UF2 规范(Microsoft UF2 spec, family 变体)打包。

用法:
  python make_uf2.py <app.bin> <out.uf2> [base_addr_hex]
  python make_uf2.py --test        # 内部往返自检(解包回来比对 payload)

UF2 块格式(512B/块): magic0 0x0A324655, magic1 0x9E5D5157,
flags 0x00002000(family), addr, len(≤476, 本工具固定 256), size, blockNo,
numBlocks, familyId=0xE48BFF56(RP2040), data[476], magicEnd 0x0AB16F30。
"""
import struct
import sys

MAGIC0 = 0x0A324655
MAGIC1 = 0x9E5D5157
MAGIC_END = 0x0AB16F30
FLAG_FAMILY = 0x00002000
FAMILY_RP2040 = 0xE48BFF56
RP2040_FLASH_BASE = 0x10000000
PAYLOAD = 256                      # RP2040 flash 页大小, 256 对齐最稳


def bin_to_uf2(data: bytes, base: int = RP2040_FLASH_BASE,
               family: int = FAMILY_RP2040) -> bytes:
    """块布局按官方 uf2conv 实现(经真机 UF2 逐块核对):
    0:magic0 4:magic1 8:flags 12:targetAddr 16:payloadSize
    20:blockNo 24:numBlocks 28:familyID 32:data[476] 508:magicEnd"""
    if len(data) % PAYLOAD:                  # 补齐到整页
        data = data + b'\x00' * (PAYLOAD - len(data) % PAYLOAD)
    total = len(data) // PAYLOAD
    out = bytearray()
    for i in range(total):
        chunk = data[i * PAYLOAD:(i + 1) * PAYLOAD]
        block = struct.pack(
            '<IIIIIIII',
            MAGIC0, MAGIC1, FLAG_FAMILY,
            base + i * PAYLOAD,               # targetAddr
            PAYLOAD,                          # payloadSize
            i, total,                         # blockNo / numBlocks
            family,                           # fileSize 槽位 = familyID
        )
        block += chunk
        block += b'\x00' * (476 - PAYLOAD)    # data 区补齐到 476
        block += struct.pack('<I', MAGIC_END)
        assert len(block) == 512
        out += block
    return bytes(out)


def uf2_to_payload(uf2: bytes):
    """解包自检用: 还原 (base, payload bytes); 非法块抛 ValueError"""
    if len(uf2) % 512:
        raise ValueError(f'uf2 长度 {len(uf2)} 不是 512 的倍数')
    blocks = [uf2[i:i + 512] for i in range(0, len(uf2), 512)]
    base = None
    total = None
    chunks = {}
    for n, b in enumerate(blocks):
        (m0, m1, flags, addr, nb, no, nt, fam) = struct.unpack_from('<8I', b)
        if m0 != MAGIC0 or m1 != MAGIC1:
            raise ValueError(f'块 {n} magic 错')
        if struct.unpack_from('<I', b, 508)[0] != MAGIC_END:
            raise ValueError(f'块 {n} 尾 magic 错')
        if flags & FLAG_FAMILY and fam != FAMILY_RP2040:
            raise ValueError(f'块 {n} family 非 RP2040: 0x{fam:08X}')
        if nb > 476:
            raise ValueError(f'块 {n} payloadSize 超限')
        if base is None:
            base, total = addr, nt
        if no != n or nt != total:
            raise ValueError(f'块 {n} 序号/总数错')
        chunks[no] = (addr, b[32:32 + nb])
    payload = bytearray()
    expect = base
    for i in range(total):
        if i not in chunks:
            raise ValueError(f'缺块 {i}')
        addr, c = chunks[i]
        if addr != expect:
            raise ValueError(f'块 {i} 地址不连续: {addr:#x} != {expect:#x}')
        payload += c
        expect += len(c)
    return base, bytes(payload)


def _selftest():
    import random
    rng = random.Random(7)
    for size in (0, 1, 255, 256, 257, 4096, 4097, 65536):
        blob = bytes(rng.randrange(256) for _ in range(size))
        if size == 0:
            continue                            # 空镜像无意义(0 块)
        uf2 = bin_to_uf2(blob, base=0x10000000)
        base, back = uf2_to_payload(uf2)
        pad = (-size) % PAYLOAD
        assert base == 0x10000000
        assert back[:size] == blob and len(back) == size + pad, \
            f'size={size} 往返不一致'
    print(f'PASS make_uf2 往返自检 ({(size + (-size) % PAYLOAD) // PAYLOAD} 块)')
    return 0


def main(argv):
    if len(argv) >= 2 and argv[1] == '--test':
        return _selftest()
    if len(argv) < 3:
        print(__doc__)
        return 2
    with open(argv[1], 'rb') as f:
        data = f.read()
    base = int(argv[3], 16) if len(argv) > 3 else RP2040_FLASH_BASE
    uf2 = bin_to_uf2(data, base=base)
    with open(argv[2], 'wb') as f:
        f.write(uf2)
    print(f'UF2: {len(uf2)} B, {len(uf2) // 512} 块, '
          f'base=0x{base:08X}, payload={len(data)} B -> {argv[2]}')
    # 写完即自检
    with open(argv[2], 'rb') as f:
        b, p = uf2_to_payload(f.read())
    assert b == base and p[:len(data)] == data
    print('PASS 写出文件解包往返一致')
    return 0


if __name__ == '__main__':
    sys.exit(main(sys.argv))