#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
merge_image.py - Merge the CAIIC bootloader and application binaries into a
                 single flash package for N32WB031KEQ6-2 (512KB).

Segments produced:
    Boot bin          -> 0x01000000 (16KB bootloader region)
    bootsetting       -> 0x01004000 (default record, generated here)
    APP bin           -> 0x01008000 (APP1 bank)

Default bootsetting record (44 bytes, little-endian, matches
Boot/src/dfu_layout.h caiic_bootsetting_t):
    magic         = 0xCA11C011
    active_bank   = 1 (APP1)
    bank1         = {start_address=0x01008000, size=<app bin size>,
                     crc32=<IEEE CRC32 of app bin>, version=0}
    bank2         = all zero
    crc32         = IEEE CRC32 over the 40 bytes above

Outputs (written to tools/out/):
    <name>.hex       Intel HEX, sparse (gaps omitted) - for pyocd/NSpyocd/J-Flash
    <name>.bin       raw binary from 0x01000000, gaps padded with 0xFF
    <name>_ota.bin   the APP image alone - this is the BLE OTA payload
    <name>_ota.json  OTA manifest for the phone app: size / crc32 / version
                     (crc32 = IEEE/zlib over the whole _ota.bin, version is
                     the value to send in OTA_BEGIN; default 0)

Usage:  python tools/merge_image.py [app_bin] [out_name] [version]
Defaults: app_bin = mcm-ddc-04/MDK-ARM/bin/embeddedSrc.bin, out_name = caiic_full,
          version = 0
Run from anywhere; paths are resolved relative to this script.
"""
import json
import os
import re
import struct
import sys
import zlib

FLASH_BASE = 0x01000000

BOOT_BIN = os.path.join("Boot", "MDK-ARM", "bin", "caiic_boot.bin")
DEFAULT_APP_BIN = os.path.join("mcm-ddc-04", "MDK-ARM", "bin", "embeddedSrc.bin")

BOOT_ADDR = 0x01000000
BOOTSETTING_ADDR = 0x01004000
APP1_ADDR = 0x01008000

BOOT_MAGIC = 0xCA11C011
ACTIVE_BANK1 = 1

OUT_DIR = "out"

HEX_REC_LEN = 16


def crc32(data):
    """IEEE CRC32 - same value as Boot/src/boot_crc32.c caiic_crc32()."""
    return zlib.crc32(data) & 0xFFFFFFFF


def default_bootsetting(app_blob):
    """Build the default 44-byte bootsetting record for APP1 = app_blob."""
    body = struct.pack(
        "<IIIIIIIIII",
        BOOT_MAGIC, ACTIVE_BANK1,
        APP1_ADDR, len(app_blob), crc32(app_blob), 0,   # bank1
        0, 0, 0, 0,                                     # bank2 (empty)
    )
    return body + struct.pack("<I", crc32(body))


def hex_record(addr, rectype, data):
    """Build one Intel HEX record string."""
    body = [len(data), (addr >> 8) & 0xFF, addr & 0xFF, rectype] + list(data)
    cks = (-sum(body)) & 0xFF
    return ":" + "".join("%02X" % b for b in body) + "%02X" % cks + "\n"


def write_hex(path, segments):
    """Write segments [(base_addr, bytes)] as Intel HEX with ELA records."""
    with open(path, "w") as f:
        ela = None
        for base, blob in segments:
            for off in range(0, len(blob), HEX_REC_LEN):
                addr = base + off
                new_ela = addr >> 16
                if new_ela != ela:
                    ela = new_ela
                    f.write(hex_record(0, 4, [(ela >> 8) & 0xFF, ela & 0xFF]))
                f.write(hex_record(addr & 0xFFFF, 0, blob[off:off + HEX_REC_LEN]))
        f.write(":00000001FF\n")


def load_bin(path, desc):
    if not os.path.isfile(path):
        print("ERROR: %s not found (%s) - build the Keil project first" % (path, desc))
        sys.exit(1)
    with open(path, "rb") as f:
        return f.read()


def read_app_fw_version(root):
    """Extract APP_FW_VERSION_NUM from the APP project's app_version.h.
    Returns None when not found."""
    hdr = os.path.join(root, "mcm-ddc-ble", "inc", "app_version.h")
    if not os.path.isfile(hdr):
        return None
    m = re.search(r"#define\s+APP_FW_VERSION_NUM\s+(0x[0-9a-fA-F]+|\d+)u?",
                open(hdr, encoding="utf-8").read())
    return int(m.group(1), 0) if m else None


def main():
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

    app_rel = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_APP_BIN
    out_name = sys.argv[2] if len(sys.argv) > 2 else "caiic_full"
    # version: explicit argv wins, otherwise take APP_FW_VERSION_NUM from the
    # firmware source so the OTA manifest always matches the flashed image
    if len(sys.argv) > 3:
        version = int(sys.argv[3], 0)
    else:
        version = read_app_fw_version(root) or 0

    boot_blob = load_bin(os.path.join(root, BOOT_BIN), "bootloader")
    app_path = app_rel if os.path.isabs(app_rel) else os.path.join(root, app_rel)
    app_blob = load_bin(app_path, "APP")
    bs_blob = default_bootsetting(app_blob)

    segments = [
        (BOOT_ADDR, boot_blob),
        (BOOTSETTING_ADDR, bs_blob),
        (APP1_ADDR, app_blob),
    ]
    for addr, blob in segments:
        print("0x%08X  (%d bytes)" % (addr, len(blob)))
    print("bootsetting: active=APP1 start=0x%08X size=%d crc32=0x%08X"
          % (APP1_ADDR, len(app_blob), crc32(app_blob)))

    out_dir = os.path.join(root, "tools", OUT_DIR)
    if not os.path.isdir(out_dir):
        os.makedirs(out_dir)

    hex_path = os.path.join(out_dir, out_name + ".hex")
    write_hex(hex_path, segments)

    # merged bin: FLASH_BASE .. end of last segment, gaps padded with 0xFF
    end = max(addr + len(blob) for addr, blob in segments)
    merged = bytearray(b"\xFF" * (end - FLASH_BASE))
    for addr, blob in segments:
        merged[addr - FLASH_BASE: addr - FLASH_BASE + len(blob)] = blob
    bin_path = os.path.join(out_dir, out_name + ".bin")
    with open(bin_path, "wb") as f:
        f.write(bytes(merged))

    # OTA payload: the APP image alone + a manifest for the phone app
    ota_bin_path = os.path.join(out_dir, out_name + "_ota.bin")
    with open(ota_bin_path, "wb") as f:
        f.write(app_blob)
    ota_json_path = os.path.join(out_dir, out_name + "_ota.json")
    with open(ota_json_path, "w") as f:
        json.dump({
            "file": os.path.basename(ota_bin_path),
            "size": len(app_blob),
            "crc32": "0x%08X" % crc32(app_blob),
            "version": version,
        }, f, indent=2)

    print("")
    print("package written:")
    print("  %s" % hex_path)
    print("  %s  (base 0x%08X, %d bytes)" % (bin_path, FLASH_BASE, len(merged)))
    print("OTA image written:")
    print("  %s  (%d bytes, crc32=0x%08X)" % (ota_bin_path, len(app_blob), crc32(app_blob)))
    print("  %s" % ota_json_path)
    return 0


if __name__ == "__main__":
    sys.exit(main())