#!/usr/bin/env python3
from __future__ import annotations

import json
import hashlib
import shutil
import subprocess
import sys
from pathlib import Path


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> None:
    if len(sys.argv) != 4:
        raise SystemExit("usage: package-firmware.py BUILD_DIR OUTPUT_DIR BOARD")

    build_dir = Path(sys.argv[1]).resolve()
    output_dir = Path(sys.argv[2]).resolve()
    board = sys.argv[3]
    repo_dir = Path(__file__).resolve().parents[1]
    version = (repo_dir / "VERSION").read_text(encoding="utf-8").strip()
    flasher = json.loads((build_dir / "flasher_args.json").read_text(encoding="utf-8"))
    flash_files: dict[str, str] = flasher["flash_files"]
    settings: dict[str, str] = flasher["flash_settings"]

    output_dir.mkdir(parents=True, exist_ok=True)
    standard_names: dict[str, str] = {}
    images: list[dict[str, str]] = []
    for offset, relative_name in sorted(flash_files.items(), key=lambda item: int(item[0], 0)):
        source = build_dir / relative_name
        if relative_name.endswith("bootloader.bin"):
            target_name = "bootloader.bin"
        elif relative_name.endswith("partition-table.bin"):
            target_name = "partition-table.bin"
        elif relative_name.endswith("device_lab.bin"):
            target_name = "firmware.bin"
        else:
            target_name = Path(relative_name).name
        shutil.copy2(source, output_dir / target_name)
        standard_names[offset] = target_name
        target = output_dir / target_name
        images.append(
            {
                "offset": offset,
                "file": target_name,
                "size": target.stat().st_size,
                "sha256": sha256(target),
            }
        )

    merge_command = [
        sys.executable,
        "-m",
        "esptool",
        "--chip",
        "esp32s3",
        "merge_bin",
        "-o",
        str(output_dir / "merged-flash.bin"),
        "--flash_mode",
        settings["flash_mode"],
        "--flash_freq",
        settings["flash_freq"],
        "--flash_size",
        settings["flash_size"],
    ]
    for offset, target_name in sorted(standard_names.items(), key=lambda item: int(item[0], 0)):
        merge_command.extend((offset, str(output_dir / target_name)))
    subprocess.run(merge_command, check=True)

    merged_image = output_dir / "merged-flash.bin"
    metadata = {
        "board": board,
        "version": version,
        "chip": "esp32s3",
        "flash_size": settings["flash_size"],
        "flash_mode": settings["flash_mode"],
        "flash_frequency": settings["flash_freq"],
        "images": images,
        "merged_image": "merged-flash.bin",
        "merged_image_size": merged_image.stat().st_size,
        "merged_image_sha256": sha256(merged_image),
        "flash_command": "python3 -m esptool --chip esp32s3 --port <PORT> write_flash 0x0 merged-flash.bin",
    }
    (output_dir / "flash.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
    (output_dir / "VERSION").write_text(version + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()