#!/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 (single-bank mode):
    Boot bin          -> 0x01000000 (16KB bootloader region)
    bootsetting       -> 0x01004000 (default record, generated here)
    APP bin           -> 0x01008000 (APP1 bank)

Dual-bank mode (--app2, optionally --with-appdata), used by
make_dual_package.bat for the appsw bank-switch test:
    Boot bin          -> 0x01000000
    bootsetting       -> 0x01004000 (both bank records filled)
    APP_DATA          -> 0x01006000 (default caiic_params_t record, 52B)
    APP1 bin          -> 0x01008000 (linked at APP1 base, 112KB bank)
    APP2 bin          -> 0x01024000 (linked at APP2 base, Keil target "APP2")

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=<app1 size>,
                     crc32=<IEEE CRC32 of app1 bin>, version=<version>}
    bank2         = same layout for the APP2 image (zeros in single-bank mode)
    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 APP1 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]
                [--app2 app2_bin [--app2-version N]] [--with-appdata]
                [--ota-app2 app2_linked_bin]
Defaults: app_bin = mcm-ddc-04/MDK-ARM/bin/embeddedSrc.bin, out_name = caiic_full,
          version = 0
--app2 merges the APP2-linked image into the package; --ota-app2 only emits
an extra OTA payload <name>_ota_app2.bin/.json (target_bank=2) without
merging. OTA manifests carry "target_bank" - the client must pick the payload
matching the bank the device is NOT running (info item 0x07 CUR_BANK).
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
APPDATA_ADDR = 0x01006000
APP1_ADDR = 0x01008000
APP2_ADDR = 0x01024000

BOOT_MAGIC = 0xCA11C011
ACTIVE_BANK1 = 1

# APP_DATA default record (mcm-ddc-ble/inc/app_params.h caiic_params_t)
PARAMS_MAGIC = 0xCA12DA7A
PARAMS_LAYOUT_VER = 1
PARAMS_LED1_BLINK_MS_DEF = 500

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, app2_blob=None, version=0, app2_version=0):
    """Build the default 44-byte bootsetting record; bank2 is zero-filled
    unless an APP2 image is given."""
    bank2 = (APP2_ADDR, len(app2_blob), crc32(app2_blob), app2_version) \
        if app2_blob else (0, 0, 0, 0)
    body = struct.pack(
        "<IIIIIIIIII",
        BOOT_MAGIC, ACTIVE_BANK1,
        APP1_ADDR, len(app_blob), crc32(app_blob), version,   # bank1
        bank2[0], bank2[1], bank2[2], bank2[3],               # bank2
    )
    return body + struct.pack("<I", crc32(body))


def default_appdata():
    """Build the 52-byte default APP_DATA parameter record
    (caiic_params_t: magic/layout_ver/led/flags/pwm/reserved[8]/crc32)."""
    body = struct.pack(
        "<IIHHI8I",
        PARAMS_MAGIC, PARAMS_LAYOUT_VER,
        PARAMS_LED1_BLINK_MS_DEF, 0,   # led1_blink_ms, flags
        0,                             # pwm_duty_pct
        0, 0, 0, 0, 0, 0, 0, 0,        # reserved[8]
    )
    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__)))

    opts = sys.argv[1:]

    app2_rel = None
    ota_app2_rel = None
    app2_version = None
    with_appdata = "--with-appdata" in opts
    if "--app2" in opts:
        i = opts.index("--app2")
        app2_rel = opts[i + 1]
    if "--ota-app2" in opts:
        i = opts.index("--ota-app2")
        ota_app2_rel = opts[i + 1]
    if "--app2-version" in opts:
        i = opts.index("--app2-version")
        app2_version = int(opts[i + 1], 0)

    # positional args minus the values consumed by the --flags
    pos = []
    skip_next = False
    for a in sys.argv[1:]:
        if skip_next:
            skip_next = False
            continue
        if a in ("--app2", "--app2-version", "--ota-app2"):
            skip_next = True
            continue
        if a == "--with-appdata":
            continue
        pos.append(a)

    app_rel = pos[0] if len(pos) > 0 else DEFAULT_APP_BIN
    out_name = pos[1] if len(pos) > 1 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(pos) > 2:
        version = int(pos[2], 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, "APP1")

    app2_blob = None
    if app2_rel:
        app2_path = app2_rel if os.path.isabs(app2_rel) else os.path.join(root, app2_rel)
        app2_blob = load_bin(app2_path, "APP2")
    if app2_version is None:
        app2_version = version   # APP2 target builds the same release

    # extra OTA payload linked for the APP2 bank (not merged into the image)
    ota_app2_blob = None
    if ota_app2_rel:
        ota_app2_path = ota_app2_rel if os.path.isabs(ota_app2_rel) \
            else os.path.join(root, ota_app2_rel)
        ota_app2_blob = load_bin(ota_app2_path, "OTA APP2 payload")
    elif app2_blob:
        ota_app2_blob = app2_blob

    bs_blob = default_bootsetting(app_blob, app2_blob, version, app2_version)

    segments = [
        (BOOT_ADDR, boot_blob),
        (BOOTSETTING_ADDR, bs_blob),
    ]
    if with_appdata:
        segments.append((APPDATA_ADDR, default_appdata()))
    segments.append((APP1_ADDR, app_blob))
    if app2_blob:
        segments.append((APP2_ADDR, app2_blob))

    for addr, blob in segments:
        print("0x%08X  (%d bytes)" % (addr, len(blob)))
    print("bootsetting: active=APP1")
    print("  bank1: start=0x%08X size=%d crc32=0x%08X ver=0x%08X"
          % (APP1_ADDR, len(app_blob), crc32(app_blob), version))
    if app2_blob:
        print("  bank2: start=0x%08X size=%d crc32=0x%08X ver=0x%08X"
              % (APP2_ADDR, len(app2_blob), crc32(app2_blob), app2_version))

    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),
            "target_bank": 1,
            "size": len(app_blob),
            "crc32": "0x%08X" % crc32(app_blob),
            "version": version,
        }, f, indent=2)

    # second OTA payload linked for the APP2 bank (when built)
    ota_app2_bin_path = None
    if ota_app2_blob:
        ota_app2_bin_path = os.path.join(out_dir, out_name + "_ota_app2.bin")
        with open(ota_app2_bin_path, "wb") as f:
            f.write(ota_app2_blob)
        with open(os.path.join(out_dir, out_name + "_ota_app2.json"), "w") as f:
            json.dump({
                "file": os.path.basename(ota_app2_bin_path),
                "target_bank": 2,
                "size": len(ota_app2_blob),
                "crc32": "0x%08X" % crc32(ota_app2_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, target bank1)" % (ota_bin_path, len(app_blob), crc32(app_blob)))
    print("  %s" % ota_json_path)
    if ota_app2_bin_path:
        print("  %s  (%d bytes, crc32=0x%08X, target bank2)"
              % (ota_app2_bin_path, len(ota_app2_blob), crc32(ota_app2_blob)))
    return 0


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