#!/usr/bin/env python3
"""
合成已加密的固件和 IUM 为最终 programmer 固件

固件布局:
  0x0000 - 0x4000:  Base firmware (16KB)
  0x4000 - 0x1DC00: 加密后的用户固件 (最大 119KB)
  0x1DC00 - 0x1DD00: IUM storage (256 bytes)
  0x1DD00 - 0x1DD10: Metadata (16 bytes)

IUM Storage 结构 (256 bytes):
  Bytes 0-252:   Modified IUM (253 bytes)
  Byte 253:      CRC8
  Bytes 254-255: CRC16 (little-endian)

用法:
  python3 merge_encrypted_firmware.py base.bin encrypted.bin ium_storage.bin -o output.bin
"""

import argparse
import struct
import sys
from pathlib import Path

# 常量定义
BASE_FIRMWARE_SIZE = 0x4000        # 16KB
USER_FIRMWARE_MAX_SIZE = 0x1DC00   # 119KB
IUM_STORAGE_SIZE = 0x100           # 256 bytes
METADATA_SIZE = 0x10               # 16 bytes

BASE_FIRMWARE_OFFSET = 0x0000
USER_FIRMWARE_OFFSET = 0x4000
IUM_OFFSET = 0x1DC00
METADATA_OFFSET = 0x1DD00


def main():
    parser = argparse.ArgumentParser(
        description='合成加密固件和IUM为最终programmer固件',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
示例:
  %(prog)s base.bin encrypted_binary.bin ium.bin -o programmer.bin

输入文件说明:
  base_firmware     - Base firmware BIN 文件 (最大 16KB)
  encrypted_binary  - 已加密的固件 BIN 文件 (*_encrypted_binary.bin)
  ium_storage       - IUM storage 文件 (256 bytes, *_ium.bin)
        '''
    )
    parser.add_argument('base_firmware', help='Base firmware BIN 文件 (最大 16KB)')
    parser.add_argument('encrypted_binary', help='已加密的固件 BIN 文件')
    parser.add_argument('ium_storage', help='IUM storage 文件 (256 bytes)')
    parser.add_argument('-o', '--output', required=True, help='输出文件路径')

    args = parser.parse_args()

    # 检查文件是否存在
    for filepath, name in [
        (args.base_firmware, 'Base firmware'),
        (args.encrypted_binary, 'Encrypted binary'),
        (args.ium_storage, 'IUM storage')
    ]:
        if not Path(filepath).exists():
            print(f"错误: {name} 文件不存在: {filepath}")
            sys.exit(1)

    # 读取文件
    base_firmware = Path(args.base_firmware).read_bytes()
    encrypted_binary = Path(args.encrypted_binary).read_bytes()
    ium_storage = Path(args.ium_storage).read_bytes()

    # 验证大小
    if len(base_firmware) > BASE_FIRMWARE_SIZE:
        print(f"错误: Base firmware 最大 {BASE_FIRMWARE_SIZE} bytes (16KB), 实际 {len(base_firmware)} bytes")
        sys.exit(1)

    if len(encrypted_binary) > USER_FIRMWARE_MAX_SIZE:
        print(f"错误: 加密固件最大 {USER_FIRMWARE_MAX_SIZE} bytes (119KB), 实际 {len(encrypted_binary)} bytes")
        sys.exit(1)

    if len(ium_storage) != IUM_STORAGE_SIZE:
        print(f"错误: IUM storage 必须是 {IUM_STORAGE_SIZE} bytes (256), 实际 {len(ium_storage)} bytes")
        sys.exit(1)

    print("合成 Programmer 固件:")
    print(f"  Base firmware:    {args.base_firmware} ({len(base_firmware)} bytes)")
    print(f"  Encrypted binary: {args.encrypted_binary} ({len(encrypted_binary)} bytes)")
    print(f"  IUM storage:      {args.ium_storage} ({len(ium_storage)} bytes)")
    print()

    # 创建输出缓冲区 (填充 0xFF)
    total_size = METADATA_OFFSET + METADATA_SIZE
    output = bytearray([0xFF] * total_size)

    # 复制 base firmware (0x0000)
    output[BASE_FIRMWARE_OFFSET:BASE_FIRMWARE_OFFSET + len(base_firmware)] = base_firmware

    # 复制加密固件 (0x4000)
    output[USER_FIRMWARE_OFFSET:USER_FIRMWARE_OFFSET + len(encrypted_binary)] = encrypted_binary

    # 复制 IUM storage (0x1DC00)
    output[IUM_OFFSET:IUM_OFFSET + IUM_STORAGE_SIZE] = ium_storage

    # 生成 metadata
    blocks = len(encrypted_binary) // 256
    last_block_size = len(encrypted_binary) % 256

    # Metadata 结构: uint16_t blocks, uint8_t last_block_size, uint8_t reserved[13]
    metadata = struct.pack('<HB', blocks, last_block_size) + bytes(13)
    output[METADATA_OFFSET:METADATA_OFFSET + METADATA_SIZE] = metadata

    # 写入输出文件
    Path(args.output).write_bytes(output)

    # 解析 IUM storage 中的 CRC
    ium_crc8 = ium_storage[253]
    ium_crc16 = ium_storage[254] | (ium_storage[255] << 8)

    print("固件结构:")
    print(f"  Base firmware:    0x{BASE_FIRMWARE_OFFSET:04X} - 0x{BASE_FIRMWARE_OFFSET + len(base_firmware):04X} ({len(base_firmware)} bytes)")
    print(f"  Encrypted binary: 0x{USER_FIRMWARE_OFFSET:04X} - 0x{USER_FIRMWARE_OFFSET + len(encrypted_binary):04X} ({len(encrypted_binary)} bytes)")
    print(f"  IUM storage:      0x{IUM_OFFSET:04X} - 0x{IUM_OFFSET + IUM_STORAGE_SIZE:04X} ({IUM_STORAGE_SIZE} bytes)")
    print(f"  Metadata:         0x{METADATA_OFFSET:04X} - 0x{METADATA_OFFSET + METADATA_SIZE:04X} ({METADATA_SIZE} bytes)")
    print()
    print("Metadata 信息:")
    print(f"  Blocks:           {blocks}")
    print(f"  Last block size:  {last_block_size} bytes")
    print(f"  Total size:       {blocks * 256 + last_block_size} bytes")
    print()
    print("IUM 校验:")
    print(f"  CRC8:  0x{ium_crc8:02X}")
    print(f"  CRC16: 0x{ium_crc16:04X}")
    print()
    print(f"✓ 输出文件: {args.output} ({len(output)} bytes)")


if __name__ == '__main__':
    main()