#!/usr/bin/env bash
#
# wrap-firmware.sh — turn a Nerves `.fw` into a bootable OpenWRT One FIT image
#
# Background: a Nerves `.fw` file contains a squashfs rootfs that
# combines the Nerves system AND your app's Erlang release at
# /srv/erlang. The OpenWRT One boots from a U-Boot FIT image
# (`openwrt-one-initramfs.itb`) containing kernel + DTB + cpio.gz
# initramfs, so we have to convert squashfs -> cpio.gz here. We
# can't use the build-time rootfs.cpio.gz from the system because
# that one doesn't have your release in it.
#
#   1. Extract data/Image, data/<dtb>, data/rootfs.img from the .fw
#   2. unsquashfs the rootfs into a temp dir
#   3. Add the bits the kernel initramfs needs that the squashfs lacks:
#         - /init shell script that mounts devtmpfs and execs /sbin/init
#   4. Repack as cpio.gz with --owner=root:root so the output is
#      reproducible and root-owned regardless of who unsquashed.
#   5. mkimage everything into a FIT (.itb) using the system's openwrt-one.its
#
# This script intentionally does NOT use sudo. Earlier versions
# preserved the original file owners with sudo unsquashfs +
# sudo mknod /dev/console, but that broke on macOS where mix-driven
# uploads can't get a TTY for sudo to prompt for a password. Without
# sudo: cpio's --owner=root:root flag forces every entry to root:root
# in the output regardless of disk owners; /dev/console is created by
# devtmpfs (which /init mounts before any console redirects) so we
# don't need to ship a static node.
#
# Usage:
#   ./scripts/wrap-firmware.sh <input.fw> <output.itb> [output.ubi]
#
# If <output.ubi> is given, the script also produces a multi-volume UBI
# image suitable for `mtd write` to /dev/mtd5 from a running Linux
# (or `mtd write spi-nand0 ... 0x100000 ...` from U-Boot). The UBI
# layout matches the OpenWrt 24.10 factory: ubootenv/ubootenv2/fip/fit/
# rootfs_data, where `fit` is the .itb we just built and `fip` is the
# pre-built ATF blob from prebuilt/openwrt-one-fip.bin.
#
# Requires (host): unzip, unsquashfs, cpio, gzip, mkimage, dtc, ubinize.

set -euo pipefail

if [ $# -lt 2 ] || [ $# -gt 3 ]; then
    echo "Usage: $0 <input.fw> <output.itb> [output.ubi]" >&2
    exit 1
fi

FW_FILE="$1"
OUT_ITB="$2"
OUT_UBI="${3:-}"

if [ ! -f "$FW_FILE" ]; then
    echo "ERROR: $FW_FILE does not exist" >&2
    exit 1
fi

# Locate this script's directory so we can find the .its template alongside it.
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SYSTEM_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
ITS_TEMPLATE="${SYSTEM_DIR}/openwrt-one.its"

if [ ! -f "$ITS_TEMPLATE" ]; then
    echo "ERROR: ITS template not found at $ITS_TEMPLATE" >&2
    exit 1
fi

# Tool checks. dtc is needed because mkimage shells out to it
# internally to compile the .its source. On Linux the u-boot-tools
# deb pulls dtc in as a dependency; on macOS Homebrew splits them,
# so users have to `brew install dtc` separately.
for tool in unzip unsquashfs cpio gzip dtc; do
    if ! command -v "$tool" >/dev/null 2>&1; then
        echo "ERROR: required tool '$tool' not found in PATH" >&2
        case "$tool" in
            dtc)
                echo "       macOS:  brew install dtc" >&2
                echo "       Linux:  apt install device-tree-compiler" >&2
                ;;
            unsquashfs)
                echo "       macOS:  brew install squashfs" >&2
                echo "       Linux:  apt install squashfs-tools" >&2
                ;;
        esac
        exit 1
    fi
done

# mkimage gets special treatment: prefer /usr/bin/mkimage (the distro-built
# u-boot-tools) over any nerves_system mkimage that may be earlier on PATH
# when this script is invoked from inside `mix upload`. The system one in
# Nerves' host/bin is compiled with -DMKIMAGE_DTC="" because Buildroot
# doesn't set CONFIG_MKIMAGE_DTC_PATH; that mkimage's internal dtc invoke
# becomes `system(" -I dts ...")` and the shell tries to run `-I` as a
# command, exploding this script. The distro mkimage has the right dtc
# path baked in.
if [ -x /usr/bin/mkimage ]; then
    MKIMAGE=/usr/bin/mkimage
elif command -v mkimage >/dev/null 2>&1; then
    MKIMAGE=$(command -v mkimage)
else
    echo "ERROR: required tool 'mkimage' not found in PATH" >&2
    exit 1
fi

WORK="$(mktemp -d -t openwrt-one-wrap.XXXXXX)"
trap 'rm -rf "$WORK"' EXIT

echo "==> Working in $WORK"

# 1-3: extract kernel, dtb, and the COMBINED rootfs (system + user
# release) from the .fw. The system-only build-time cpio.gz won't
# work as initramfs because erlinit can't find /srv/erlang/* without
# the user release.
echo "==> Extracting kernel, dtb, rootfs from $FW_FILE"
unzip -p "$FW_FILE" data/Image                    > "$WORK/Image"
unzip -p "$FW_FILE" data/mt7981b-openwrt-one.dtb  > "$WORK/mt7981b-openwrt-one.dtb"
unzip -p "$FW_FILE" data/rootfs.img               > "$WORK/rootfs.squashfs"

# 4: unsquashfs (no sudo — owners get reset to root:root by cpio below).
echo "==> Unpacking squashfs..."
unsquashfs -no-progress -d "$WORK/rootfs" "$WORK/rootfs.squashfs" > /dev/null

# 5: add /init. /dev/console is auto-created by devtmpfs (mounted
# inside /init); no static node needed.
echo "==> Adding initramfs /init"
cat > "$WORK/rootfs/init" <<'EOF'
#!/bin/sh
# devtmpfs does not get automounted for initramfs
/bin/mount -t devtmpfs devtmpfs /dev

# use the /dev/console device node from devtmpfs if possible (avoids
# glibc ttyname_r confusion). Wrapped in a subshell so a failing exec
# doesn't terminate the parent.
if (exec 0</dev/console) 2>/dev/null; then
    exec 0</dev/console
    exec 1>/dev/console
    exec 2>/dev/console
fi

exec /sbin/init "$@"
EOF
chmod +x "$WORK/rootfs/init"

# 6: repack as cpio.gz. --owner=0:0 forces every entry to be owned
# by uid 0 / gid 0 (root) regardless of who unsquashed (no sudo
# needed). We use numeric IDs because BSD cpio (macOS default)
# resolves --owner=root:root by looking up "root" as a group in
# /etc/group, which fails on macOS where the gid-0 group is "wheel".
# We deliberately don't pass --reproducible: it's a GNU cpio
# extension that BSD cpio rejects, and the cpio byte-reproducibility
# doesn't actually matter for booting.
echo "==> Packing cpio.gz..."
( cd "$WORK/rootfs" && \
  find . -mindepth 1 | LC_ALL=C sort | \
  cpio --owner=0:0 --quiet -o -H newc ) | gzip -9 > "$WORK/rootfs.cpio.gz"

CPIO_SIZE=$(stat -c%s "$WORK/rootfs.cpio.gz" 2>/dev/null || stat -f%z "$WORK/rootfs.cpio.gz")
echo "    rootfs.cpio.gz: $((CPIO_SIZE / 1024 / 1024)) MiB"

# 7: build the FIT image
echo "==> Building FIT image..."
cp "$ITS_TEMPLATE" "$WORK/openwrt-one.its"
( cd "$WORK" && "$MKIMAGE" -f openwrt-one.its "$OUT_ITB" >/dev/null )

OUT_SIZE=$(stat -c%s "$OUT_ITB" 2>/dev/null || stat -f%z "$OUT_ITB")
echo "==> Wrote $OUT_ITB ($((OUT_SIZE / 1024 / 1024)) MiB)"

# 8 (optional): wrap the .itb in a multi-volume UBI image for SPI NAND.
if [ -n "$OUT_UBI" ]; then
    UBINIZE=""
    if command -v ubinize >/dev/null 2>&1; then
        UBINIZE=ubinize
    else
        # A from-source Buildroot build leaves ubinize in host/sbin or
        # host/bin under the system artifact; portable (cached) artifacts
        # don't ship host tools, so users on those will need mtd-utils.
        for cand in \
            "${HOME}/.nerves/artifacts/nerves_system_openwrt_one-portable-"*/host/sbin/ubinize \
            "${HOME}/.nerves/artifacts/nerves_system_openwrt_one-portable-"*/host/bin/ubinize \
            "${SYSTEM_DIR}/.nerves/artifacts/nerves_system_openwrt_one-portable-"*/host/sbin/ubinize \
            "${SYSTEM_DIR}/.nerves/artifacts/nerves_system_openwrt_one-portable-"*/host/bin/ubinize; do
            if [ -x "$cand" ]; then
                UBINIZE="$cand"
                break
            fi
        done
    fi

    if [ -z "$UBINIZE" ]; then
        echo "ERROR: ubinize not found on host." >&2
        echo "  Linux:  sudo apt install mtd-utils" >&2
        echo "  macOS:  no Homebrew package; build mtd-utils from source or use Linux" >&2
        exit 1
    else
        UBI_CFG_TEMPLATE="${SYSTEM_DIR}/ubinize-fit.cfg"
        if [ ! -f "$UBI_CFG_TEMPLATE" ]; then
            echo "WARNING: $UBI_CFG_TEMPLATE not found, skipping UBI image generation" >&2
        else
            echo "==> Building UBI image (multi-volume layout for SPI NAND)..."
            UBI_CFG="$WORK/ubinize-fit.cfg"

            # Generate a populated U-Boot environment binary so freshly
            # flashed devices already have the Nerves metadata. Without
            # this, NervesMOTD shows "unknown" until UBootEnv.write/1 is
            # called manually -- nerves_runtime caches KV at boot so just
            # writing later doesn't help.
            #
            # We parse meta-* lines from the .fw file's meta.conf and emit
            # an env text file in the format mkenvimage expects, then turn
            # it into a 64 KiB binary block (matching /etc/fw_env.config's
            # env_size = 0x10000).
            MKENVIMAGE=""
            if command -v mkenvimage >/dev/null 2>&1; then
                MKENVIMAGE=mkenvimage
            else
                for cand in \
                    "${HOME}/.nerves/artifacts/nerves_system_openwrt_one-portable-"*/host/bin/mkenvimage \
                    "${SYSTEM_DIR}/.nerves/artifacts/nerves_system_openwrt_one-portable-"*/host/bin/mkenvimage; do
                    if [ -x "$cand" ]; then MKENVIMAGE="$cand"; break; fi
                done
            fi

            if [ -z "$MKENVIMAGE" ]; then
                echo "WARNING: mkenvimage not found. ubootenv volume will be empty," >&2
                echo "         meaning Nerves metadata won't be populated until you" >&2
                echo "         call UBootEnv.write/1 manually after first boot." >&2
                # Still need an env binary for ubinize, even if empty.
                printf '\xff%.0s' $(seq 1 126976) > "$WORK/uboot-env.bin"
            else
                echo "==> Generating U-Boot env binary from .fw metadata..."
                # Extract meta-* lines from meta.conf and convert to env format.
                # meta.conf entries look like:  meta-product=openwrt_one_test
                #                               meta-version=0.1.0
                #                               meta-author="Herman verschooten"
                # We strip surrounding quotes and rename them to nerves_fw_*.
                ENV_TXT="$WORK/uboot-env.txt"
                # Start from the OpenWrt U-Boot default env template so the
                # baked env contains a complete bootcmd / boot_* / ubi_* /
                # bootmenu_* tree. With those present U-Boot accepts our
                # env directly (CRC matches at 0x1F000), no warnings from
                # boardid, and Nerves still gets nerves_fw_* alongside.
                cat "${SYSTEM_DIR}/prebuilt/uboot-env-template.txt" > "$ENV_TXT"
                {
                    printf '\n# nerves_fw_* keys appended at build time:\n'
                    # Slot tracking: assume single-slot install on first flash.
                    printf 'nerves_fw_active=a\n'

                    # Pull each meta-* field from meta.conf, transform name.
                    unzip -p "$FW_FILE" meta.conf | awk '
                        /^meta-/ {
                            # Remove "meta-" prefix
                            sub(/^meta-/, "")
                            # Split on first "="
                            eq = index($0, "=")
                            if (eq == 0) next
                            key = substr($0, 1, eq - 1)
                            val = substr($0, eq + 1)
                            # Strip surrounding double quotes if present
                            gsub(/^"|"$/, "", val)
                            # Map fwup meta-* names to Nerves env keys
                            if (key == "product")       printf "a.nerves_fw_product=%s\n", val
                            else if (key == "version")  printf "a.nerves_fw_version=%s\n", val
                            else if (key == "platform") printf "a.nerves_fw_platform=%s\n", val
                            else if (key == "architecture") printf "a.nerves_fw_architecture=%s\n", val
                            else if (key == "author")   printf "a.nerves_fw_author=%s\n", val
                            else if (key == "description") printf "a.nerves_fw_description=%s\n", val
                            else if (key == "vcs-identifier") printf "a.nerves_fw_vcs_identifier=%s\n", val
                            else if (key == "misc")     printf "a.nerves_fw_misc=%s\n", val
                            else if (key == "uuid")     printf "a.nerves_fw_uuid=%s\n", val
                        }
                    '
                } >> "$ENV_TXT"
                echo "    appended nerves_fw_* keys:"
                grep '^nerves_fw_\|^a\.nerves_fw_' "$ENV_TXT" | sed 's/^/      /'
                # -r: redundant env (5-byte header = 4-byte CRC + 1-byte flags)
                # because we have ubootenv + ubootenv2 volumes. Without -r the
                # CRC check in fw_printenv / Erlang UBootEnv fails because
                # the data is offset by one byte.
                # env_size 0x1F000 (= one UBI LEB) MUST match CONFIG_ENV_SIZE
                # in OpenWrt's mt7981 U-Boot. We start from the OpenWrt
                # default env template (with bootcmd, boot_*, ubi_*, etc.)
                # and append nerves_fw_* keys, so the resulting env satisfies
                # both U-Boot (full boot scripts present, CRC matches) and
                # Nerves (nerves_fw_* present for KV reads).
                "$MKENVIMAGE" -r -s 0x1F000 -o "$WORK/uboot-env.bin" "$ENV_TXT"
            fi

            # The template uses BOARD_DIR and BINARIES_DIR placeholders that
            # we substitute. Point BOARD_DIR at the system tree (so it can
            # find prebuilt/openwrt-one-fip.bin) and BINARIES_DIR at $WORK
            # (where we just produced the .itb and uboot-env.bin). The .itb
            # in $WORK has the name openwrt-one-initramfs.itb because that's
            # what the .cfg references; copy our $OUT_ITB to that name.
            cp "$OUT_ITB" "$WORK/openwrt-one-initramfs.itb"
            sed -e "s|BOARD_DIR|${SYSTEM_DIR}|g" \
                -e "s|BINARIES_DIR|${WORK}|g" \
                "$UBI_CFG_TEMPLATE" > "$UBI_CFG"
            "$UBINIZE" -p 128KiB -m 2048 -o "$OUT_UBI" "$UBI_CFG"
            UBI_SIZE=$(stat -c%s "$OUT_UBI" 2>/dev/null || stat -f%z "$OUT_UBI")
            echo "==> Wrote $OUT_UBI ($((UBI_SIZE / 1024 / 1024)) MiB)"
        fi
    fi
fi

echo ""
echo "To boot via TFTP from U-Boot (no NAND changes):"
echo "  tftpboot \$loadaddr $(basename "$OUT_ITB")"
echo "  bootm"
if [ -n "$OUT_UBI" ] && [ -f "$OUT_UBI" ]; then
    echo ""
    echo "To install on NAND from a running OpenWrt (or any Linux on this board):"
    echo "  scp $(basename "$OUT_UBI") root@<ip>:/tmp/"
    echo "  ssh root@<ip>"
    echo "    ubidetach -p /dev/mtd5 2>/dev/null || true"
    echo "    flash_erase /dev/mtd5 0 0"
    echo "    nandwrite -p /dev/mtd5 /tmp/$(basename "$OUT_UBI")"
    echo "    reboot"
    echo ""
    echo "To install on a blank NAND from U-Boot:"
    echo "  tftpboot \$loadaddr $(basename "$OUT_UBI")"
    echo "  ubi detach"
    echo "  mtd erase ubi"
    echo "  mtd write spi-nand0 \$loadaddr 0x100000 \$filesize"
    echo "  reset"
fi