mesh-forge/scripts/emit-flash-manifest.py

#!/usr/bin/env python3
"""
Emit flash-manifest.json for Mesh Forge USB flasher.

If BUILD_DIR/flash-manifest.json already exists (project-supplied), merge in
targetFamily from PlatformIO when missing.

Otherwise, if Meshtastic-style *.mt.json is present, synthesize a manifest
from partition table + on-disk artifacts.

Optional args: PROJECT_ROOT TARGET_ENV — merged PIO config for that env fills
targetFamily (and platform/board) for the USB flasher UI.
"""
from __future__ import annotations

import glob
import json
import os
import re
import sys


def _normalize_platform(platform: str | None) -> str:
    if not platform:
        return ""
    p = platform.strip().lower()
    if "#" in p:
        p = p.split("#", 1)[0].strip()
    if "@" in p:
        p = p.split("@", 1)[0].strip()
    if "/" in p:
        p = p.rsplit("/", 1)[-1].strip()
    return p


def _platform_to_target_family(platform: str | None, board: str | None) -> str:
    pl = _normalize_platform(platform)
    b = (board or "").lower()
    if "8266" in pl or "esp8266" in b:
        return "esp8266"
    if "nrf52" in pl or "nrf52840" in b or "nrf52833" in b or "nrf52832" in b:
        return "nrf52"
    if "rp2040" in pl or "rp2040" in b or "raspberrypi" in pl:
        return "rp2040"
    if "espressif32" in pl or "esp32" in pl or "esp32" in b or "esp32c3" in b or "esp32s3" in b:
        return "esp32"
    return "unknown"


def _resolve_pio_target_family(project_root: str, target_env: str) -> tuple[str, str | None, str | None]:
    try:
        from platformio.project.config import ProjectConfig
    except ImportError:
        print("platformio not installed; targetFamily will be unknown", file=sys.stderr)
        return "unknown", None, None

    ini = os.path.join(project_root, "platformio.ini")
    if not os.path.isfile(ini):
        return "unknown", None, None

    old = os.getcwd()
    try:
        os.chdir(project_root)
        config = ProjectConfig("platformio.ini")
        section = f"env:{target_env}"
        if section not in config.sections():
            print(f"PIO env not found: {target_env!r}", file=sys.stderr)
            return "unknown", None, None
        platform = config.get(section, "platform")
        board = config.get(section, "board")
        fam = _platform_to_target_family(platform, board)
        return fam, platform, board
    except Exception as e:
        print(f"PIO targetFamily resolution failed: {e}", file=sys.stderr)
        return "unknown", None, None
    finally:
        os.chdir(old)


def _merge_target_family_meta(
    manifest: dict,
    target_family: str,
    platform: str | None,
    board: str | None,
) -> dict:
    """Add targetFamily / platform / board only when absent."""
    out = dict(manifest)
    if "targetFamily" not in out:
        out["targetFamily"] = target_family
    if platform and "platform" not in out:
        out["platform"] = platform.strip() if isinstance(platform, str) else platform
    if board and "board" not in out:
        out["board"] = board.strip() if isinstance(board, str) else board
    return out


def _write_manifest(path: str, data: dict) -> None:
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)


def parse_offset(v: object) -> int | None:
    if v is None:
        return None
    if isinstance(v, bool):
        return None
    if isinstance(v, int):
        return v
    s = str(v).strip()
    if not s:
        return None
    if s.lower().startswith("0x"):
        return int(s, 16)
    return int(s, 10)


def list_basenames(build_dir: str) -> set[str]:
    out: set[str] = set()
    for p in glob.glob(os.path.join(build_dir, "*")):
        if os.path.isfile(p):
            out.add(os.path.basename(p))
    return out


def part_offset_for_slot(parts: list[dict], part_name: str) -> int | None:
    for p in parts:
        if str(p.get("name", "")) == part_name:
            o = parse_offset(p.get("offset"))
            if o is not None:
                return o
    for p in parts:
        if str(p.get("subtype", "")) == part_name:
            o = parse_offset(p.get("offset"))
            if o is not None:
                return o
    if part_name == "app0":
        for p in parts:
            if p.get("type") == "app" and p.get("subtype") == "ota_0":
                o = parse_offset(p.get("offset"))
                if o is not None:
                    return o
    if part_name == "spiffs":
        for p in parts:
            if str(p.get("subtype", "")) == "spiffs":
                o = parse_offset(p.get("offset"))
                if o is not None:
                    return o
    return None


def factory_app_offset(parts: list[dict]) -> int:
    for p in parts:
        if str(p.get("subtype", "")) == "factory":
            o = parse_offset(p.get("offset"))
            if o is not None:
                return o
    for p in parts:
        if p.get("type") == "app" and p.get("subtype") == "ota_0":
            o = parse_offset(p.get("offset"))
            if o is not None:
                return o
    return 0x10000


def ota1_offset(parts: list[dict]) -> int | None:
    for p in parts:
        if str(p.get("subtype", "")) == "ota_1":
            return parse_offset(p.get("offset"))
    return None


def spiffs_offset(parts: list[dict]) -> int | None:
    return part_offset_for_slot(parts, "spiffs")


def _offset_int(im: dict) -> int:
    o = im["offset"]
    return int(o) if isinstance(o, int) else parse_offset(o) or 0


def _dedupe_same_offset(images: list[dict]) -> list[dict]:
    """
    One esptool image per physical offset. Meshtastic maps both firmware-*.bin (app0)
    and firmware-*.factory.bin to the same ota_0 slot — keep factory, drop the duplicate.
    """
    buckets: dict[int, list[dict]] = {}
    for im in images:
        buckets.setdefault(_offset_int(im), []).append(im)

    def rank(im: dict) -> tuple:
        f = im["file"]
        opt = im.get("optional") is True
        if re.search(r"\.factory\.bin$", f, re.I):
            return (0, 0 if not opt else 1, f)
        if not opt:
            return (1, 0, f)
        return (2, 0, f)

    out: list[dict] = []
    for off in sorted(buckets):
        group = buckets[off]
        out.append(group[0] if len(group) == 1 else min(group, key=rank))
    return out


def emit_from_mt(build_dir: str, mt: dict) -> dict | None:
    parts: list[dict] = mt.get("part") or []
    if not parts:
        return None

    names = list_basenames(build_dir)
    images: list[dict] = []
    seen: set[str] = set()

    def add(file: str, offset: int, optional: bool = False) -> None:
        if file not in names or file in seen:
            return
        row: dict = {"file": file, "offset": offset}
        if optional:
            row["optional"] = True
        images.append(row)
        seen.add(file)

    add("bootloader.bin", 0x1000)
    add("partitions.bin", 0x8000)
    add("boot_app0.bin", 0xE000)

    for entry in mt.get("files") or []:
        fname = entry.get("name")
        part_name = entry.get("part_name")
        if not fname or not part_name or fname not in names:
            continue
        off = part_offset_for_slot(parts, str(part_name))
        if off is None:
            continue
        opt = bool(
            fname.startswith("littlefs-")
            or re.match(r"^mt-.+-ota\.bin$", fname, re.I)
            or (
                re.match(r"^firmware-.+\.bin$", fname, re.I)
                and not re.search(r"\.factory\.bin$", fname, re.I)
            )
        )
        add(fname, off, optional=opt)

    factory_bins = sorted(n for n in names if re.match(r"^firmware-.+\.factory\.bin$", n, re.I))
    if factory_bins:
        add(factory_bins[0], factory_app_offset(parts))

    for n in sorted(names):
        if n.startswith("littlefs-") and n.endswith(".bin"):
            off = spiffs_offset(parts)
            if off is not None:
                add(n, off, optional=True)

    for n in sorted(names):
        if re.match(r"^mt-.+-ota\.bin$", n, re.I):
            off = ota1_offset(parts)
            if off is not None:
                add(n, off, optional=True)

    if not images:
        return None

    images = _dedupe_same_offset(images)

    if not any(re.match(r"^firmware-.+\.bin$", im["file"], re.I) for im in images):
        return None

    images.sort(key=_offset_int)
    return {"images": images}


def pick_mt_json(build_dir: str) -> str | None:
    paths = glob.glob(os.path.join(build_dir, "*.mt.json"))
    if not paths:
        return None
    if len(paths) == 1:
        return paths[0]
    fw = [p for p in paths if re.search(r"firmware-.+\.mt\.json$", os.path.basename(p), re.I)]
    return fw[0] if fw else paths[0]


def main() -> int:
    if len(sys.argv) < 2:
        print(
            "usage: emit-flash-manifest.py BUILD_DIR [PROJECT_ROOT TARGET_ENV]",
            file=sys.stderr,
        )
        return 2
    build_dir = os.path.abspath(sys.argv[1])
    project_root: str | None = None
    target_env: str | None = None
    if len(sys.argv) >= 4:
        project_root = os.path.abspath(sys.argv[2])
        target_env = sys.argv[3]

    pio_family, pio_platform, pio_board = (
        _resolve_pio_target_family(project_root, target_env)
        if project_root and target_env
        else ("unknown", None, None)
    )

    out_path = os.path.join(build_dir, "flash-manifest.json")

    if os.path.isfile(out_path):
        with open(out_path, encoding="utf-8") as f:
            manifest = json.load(f)
        if not isinstance(manifest, dict) or not isinstance(manifest.get("images"), list):
            print(f"Invalid existing {out_path}", file=sys.stderr)
            return 1
        merged = _merge_target_family_meta(manifest, pio_family, pio_platform, pio_board)
        if merged != manifest:
            _write_manifest(out_path, merged)
            print(f"Merged targetFamily into {out_path}")
        else:
            print(f"Keeping existing {out_path} (targetFamily already set)")
        return 0

    mt_path = pick_mt_json(build_dir)
    if not mt_path:
        print("No *.mt.json; not emitting flash-manifest.json")
        return 0

    with open(mt_path, encoding="utf-8") as f:
        mt = json.load(f)

    manifest = emit_from_mt(build_dir, mt)
    if not manifest:
        print("Could not synthesize flash-manifest.json from mt.json")
        return 0

    manifest = _merge_target_family_meta(manifest, pio_family, pio_platform, pio_board)
    _write_manifest(out_path, manifest)
    print(f"Wrote {out_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())