#!/usr/bin/env python3
"""
PlatformIO extra script: generate DFU .zip after build.

Add to platformio.ini:
  extra_scripts = post:flasher/pio_create_dfu_zip.py
"""

import sys
import os
import json
import struct
import zipfile

# Add flasher dir to path for importing the main module
sys.path.insert(0, os.path.join(os.path.dirname(__file__)))

try:
    from create_dfu_zip import hex_to_bin, uf2_to_bin, create_dfu_zip
except ImportError:
    # Fallback: define minimal conversion inline
    def hex_to_bin(hex_path):
        from intelhex import IntelHex
        ih = IntelHex(hex_path)
        min_addr = ih.minaddr() or 0
        max_addr = ih.maxaddr() or 0
        size = max_addr - min_addr + 1
        page_size = 0x1000
        aligned = ((size + page_size - 1) // page_size) * page_size
        return bytes(ih.tobinarray(start=min_addr, size=aligned)), min_addr

    def create_dfu_zip(bin_data, base_addr, fw_version=1, hw_version=52):
        init_packet = bytearray()
        init_packet.append(0x01)
        init_packet.append(0x04)
        init_packet += struct.pack('<I', 0xFFFE)
        init_packet += struct.pack('<I', 0xFFFFFFFF)
        init_packet += struct.pack('<I', hw_version)
        init_packet += struct.pack('<I', fw_version)
        init_packet += struct.pack('<II', 0, 0)
        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 create_dfu_zip_from_build(target, source, env):
    """PlatformIO post-build hook."""
    build_dir = env.subst("$BUILD_DIR")
    prog_name = env.subst("$PROGNAME")
    hex_file = os.path.join(build_dir, prog_name + ".hex")

    if not os.path.isfile(hex_file):
        print(f"  [DFU] {hex_file} not found, skipping")
        return

    try:
        bin_data, base_addr = hex_to_bin(hex_file)
    except Exception as e:
        print(f"  [DFU] Error: {e}")
        return

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

    # Name the zip after the build target (environment name)
    env_name = env.subst("$PIOENV")
    # Extract firmware name from env: e.g. Heltec_t114_without_display_beacon_sensor_ble -> Heltec_T114_Beacon_BLE
    parts = env_name.split("_")
    if "beacon" in parts:
        fw_name = "Heltec_T114_Beacon"
    elif "companion" in parts:
        fw_name = "Heltec_T114_Companion_Radio"
    elif "repeater" in parts:
        fw_name = "Heltec_T114_Repeater"
    elif "room" in parts:
        fw_name = "Heltec_T114_Room_Server"
    else:
        fw_name = "firmware"

    if "ble" in parts:
        fw_name += "_BLE"
    elif "usb" in parts:
        fw_name += "_USB"

    zip_name = f"{fw_name}.dfu.zip"
    zip_path = os.path.join(build_dir, zip_name)

    with zipfile.ZipFile(zip_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(zip_path)
    print(f"  [DFU] Created: {zip_path} ({size / 1024:.0f} KB)")


# Register hook
Import("env")
env.AddPostAction("$BUILD_DIR/${PROGNAME}.hex", create_dfu_zip_from_build)

# Also trigger after .uf2 creation if that action exists
try:
    env.AddPostAction("$BUILD_DIR/${PROGNAME}.uf2", create_dfu_zip_from_build)
except:
    pass