#!/usr/bin/env python3

import json
import re
import shutil
from pathlib import Path


BUILD_ARTIFACTS = (
    "bootloader.bin",
    "partitions.bin",
    "firmware.bin",
    "firmware.elf",
    "firmware.map",
)

MANIFEST_PARTS = (
    {"path": "bootloader.bin", "offset": 4096},
    {"path": "partitions.bin", "offset": 32768},
    {"path": "boot_app0.bin", "offset": 57344},
    {"path": "firmware.bin", "offset": 65536},
)

VERSION_PATTERN = re.compile(r'Eggduino-Firmware V([^"\\]+)')


def extract_firmware_version(header_text: str) -> str:
    match = VERSION_PATTERN.search(header_text)
    if not match:
        raise ValueError("Could not extract firmware version from include/EggDuino.h")
    return match.group(1).strip()


def build_manifest(version: str) -> dict:
    return {
        "name": "EggDuino ESP32 Firmware",
        "version": version,
        "new_install_prompt_erase": True,
        "builds": [
            {
                "chipFamily": "ESP32",
                "parts": list(MANIFEST_PARTS),
            }
        ],
    }


def package_bundle(build_dir: Path, boot_app0_path: Path, header_path: Path, output_dir: Path) -> None:
    build_dir = Path(build_dir)
    boot_app0_path = Path(boot_app0_path)
    header_path = Path(header_path)
    output_dir = Path(output_dir)

    if output_dir.exists():
        shutil.rmtree(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    for artifact_name in BUILD_ARTIFACTS:
        artifact_path = build_dir / artifact_name
        if not artifact_path.is_file():
            raise FileNotFoundError(f"Missing build artifact: {artifact_path}")
        shutil.copy2(artifact_path, output_dir / artifact_name)

    if not boot_app0_path.is_file():
        raise FileNotFoundError(f"Missing boot_app0.bin: {boot_app0_path}")
    shutil.copy2(boot_app0_path, output_dir / "boot_app0.bin")

    version = extract_firmware_version(header_path.read_text(encoding="utf-8"))
    manifest_path = output_dir / "manifest.json"
    manifest_path.write_text(
        json.dumps(build_manifest(version), indent=4) + "\n",
        encoding="utf-8",
    )


def _get_boot_app0_path(env) -> Path:
    framework_dir = env.PioPlatform().get_package_dir("framework-arduinoespressif32")
    if not framework_dir:
        raise FileNotFoundError("PlatformIO framework-arduinoespressif32 package is not installed")

    boot_app0_path = Path(framework_dir) / "tools" / "partitions" / "boot_app0.bin"
    if not boot_app0_path.is_file():
        raise FileNotFoundError(f"Could not find boot_app0.bin in {boot_app0_path.parent}")
    return boot_app0_path


def _package_platformio_bundle(target, source, env) -> None:
    del target, source
    project_dir = Path(env.subst("$PROJECT_DIR"))
    package_bundle(
        build_dir=Path(env.subst("$BUILD_DIR")),
        boot_app0_path=_get_boot_app0_path(env),
        header_path=project_dir / "include" / "EggDuino.h",
        output_dir=project_dir / "firmware",
    )
    print(f"Packaged firmware bundle to {project_dir / 'firmware'}")


if "Import" in globals():
    Import("env")
    env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", _package_platformio_bundle)