#!/usr/bin/env python3
"""
Convert .hex or .uf2 to nRF52 DFU .zip for Web Serial flashing.

Usage:
  python3 create_dfu_zip.py firmware.hex firmware.dfu.zip
  python3 create_dfu_zip.py firmware.uf2 firmware.dfu.zip
"""

import sys
import os
import json
import struct
import zipfile
import hashlib

def hex_to_bin(hex_path):
    """Convert Intel HEX to raw binary."""
    from intelhex import IntelHex
    ih = IntelHex(hex_path)
    # Use the full address range or just the used portion
    min_addr = ih.minaddr() if ih.minaddr() is not None else 0
    max_addr = ih.maxaddr() if ih.maxaddr() is not None else 0
    size = max_addr - min_addr + 1
    # Align to page size (4KB for nRF52840)
    page_size = 0x1000
    aligned_size = ((size + page_size - 1) // page_size) * page_size
    data = ih.tobinarray(start=min_addr, size=aligned_size)
    return bytes(data), min_addr

def uf2_to_bin(uf2_path):
    """Convert UF2 to raw binary."""
    FAMILY_NRF52840 = 0xADA52840
    data = {}
    with open(uf2_path, 'rb') as f:
        while True:
            block = f.read(512)
            if not block or len(block) < 512:
                break
            magic0, magic1, flags, addr, size, seq, num = struct.unpack_from('<IIIIIII', block, 0)
            if magic0 != 0x0A324655 or magic1 != 0x9E5D5157:
                continue
            payload = block[32:32+size]
            for i in range(0, len(payload), 4):
                word = struct.unpack_from('<I', payload, i)[0]
                offset = addr + i
                data[offset] = word
    if not data:
        raise ValueError("No valid UF2 blocks found")
    addresses = sorted(data.keys())
    min_addr = addresses[0] & ~0xFFF  # align to 4K
    max_addr = addresses[-1]
    page_size = 0x1000
    size = max_addr - min_addr + 1
    aligned_size = ((size + page_size - 1) // page_size) * page_size
    bin_data = bytearray(aligned_size)
    for addr, word in data.items():
        offset = addr - min_addr
        if offset + 4 <= len(bin_data):
            struct.pack_into('<I', bin_data, offset, word)
    return bytes(bin_data), min_addr

def create_dfu_zip(bin_data, base_addr, fw_version=1, hw_version=52):
    """Create nRF DFU zip from binary data."""

    # Build init packet (.dat) per nRF DFU spec
    # Format (little-endian):
    #   [0x01] - signature (DFU init packet)
    #   [fw_type: 1] - 0x04 = application
    #   [sd_count: 4] - number of SD requirements (1)
    #   [sd_array: sd_count * 4] - SD requirement (0xFFFE = any)
    #   [hw_version: 4]
    #   [fw_version: 4]
    #   [fwid_type: 2] - 0x0001 = SHA-256 (not used, set to 0)
    #   [fwid_length: 2]
    #   [fwid: fwid_length]

    # Simpler init packet format as expected by dfu.js:
    # It expects init_packet_data with components.data array

    init_packet = bytearray()
    init_packet.append(0x01)           # signature
    init_packet.append(0x04)           # fw_type: application
    # softdevice requirements
    init_packet += struct.pack('<I', 0xFFFE)  # SD required: any
    init_packet += struct.pack('<I', 0xFFFFFFFF)  # terminator
    init_packet += struct.pack('<I', hw_version)   # HW version
    init_packet += struct.pack('<I', fw_version)   # FW version
    # FWID (empty)
    init_packet += struct.pack('<II', 0, 0)  # type=0, len=0

    # build manifest
    manifest = {
        "manifest": {
            "application": {
                "bin_file": "firmware.bin",
                "dat_file": "firmware.dat",
                "init_packet_data": {
                    "fw_version": fw_version,
                    "hw_version": hw_version,
                    "softdevice_req": [0xFFFE],
                    "components": [
                        {
                            "data": list(init_packet)
                        }
                    ]
                }
            }
        }
    }

    return manifest, bin_data, bytes(init_packet)

def main():
    if len(sys.argv) < 3:
        print("Usage: create_dfu_zip.py <input.hex|uf2> <output.zip> [fw_version] [hw_version]")
        sys.exit(1)

    input_path = sys.argv[1]
    output_path = sys.argv[2]
    fw_version = int(sys.argv[3]) if len(sys.argv) > 3 else 1
    hw_version = int(sys.argv[4]) if len(sys.argv) > 4 else 52

    ext = os.path.splitext(input_path)[1].lower()

    if ext == '.hex':
        try:
            bin_data, base_addr = hex_to_bin(input_path)
        except ImportError:
            print("Error: intelhex package required for .hex files.")
            print("Install: pip install intelhex")
            sys.exit(1)
    elif ext == '.uf2':
        bin_data, base_addr = uf2_to_bin(input_path)
    else:
        print(f"Unsupported format: {ext}")
        sys.exit(1)

    manifest, fw_bin, fw_dat = create_dfu_zip(bin_data, base_addr, fw_version, hw_version)

    with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("manifest.json", json.dumps(manifest, indent=2))
        zf.writestr("firmware.bin", fw_bin)
        zf.writestr("firmware.dat", fw_dat)

    size = os.path.getsize(output_path)
    print(f"DFU zip created: {output_path}")
    print(f"  Size: {size} bytes ({size/1024:.0f} KB)")
    print(f"  Firmware: {len(fw_bin)} bytes @ 0x{base_addr:08X}")

if __name__ == '__main__':
    main()