#!/usr/bin/env bash
#
# wrap-firmware.sh — turn a Nerves `.fw` into a bootable OpenWRT One FIT image
#
# Background: a Nerves `.fw` file is a zip containing a squashfs rootfs (with
# the Erlang VM and your application's release inside). The OpenWRT One boot
# chain expects a U-Boot FIT image (`openwrt-one-initramfs.itb`) containing
# kernel + DTB + cpio.gz initramfs. This script bridges the two:
#
#   1. Extract  data/rootfs.img  (squashfs)         from the .fw
#   2. Extract  data/Image                          from the .fw
#   3. Extract  data/mt7981b-openwrt-one.dtb        from the .fw
#   4. unsquashfs the rootfs into a temp dir
#   5. Add the bits the kernel initramfs needs that the squashfs lacks:
#         - /init shell script that mounts devtmpfs and execs /sbin/init
#         - /dev/console char device node (5,1)
#   6. Repack as cpio.gz
#   7. mkimage everything into a FIT (.itb) using the system's openwrt-one.its
#
# 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, ubinize,
# sudo (for mknod & preserving file ownership in the cpio).
#
# This is a Phase 2 development convenience script. Phase 3 will replace
# this with proper fwup.conf integration so `mix firmware` produces the
# .itb directly and `mix upload` can write to NAND via UBI volumes.

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
for tool in unzip unsquashfs cpio gzip sudo; do
    if ! command -v "$tool" >/dev/null 2>&1; then
        echo "ERROR: required tool '$tool' not found in PATH" >&2
        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 'sudo rm -rf "$WORK"' EXIT

echo "==> Working in $WORK"

# 1-3: extract the three things we need from the .fw
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 the rootfs (sudo to preserve owners/permissions)
echo "==> Unpacking squashfs..."
sudo unsquashfs -no-progress -d "$WORK/rootfs" "$WORK/rootfs.squashfs" > /dev/null

# 5: add the initramfs essentials. The squashfs is built without device nodes
#    or a top-level /init because Nerves normally mounts it as a real
#    block-device root, not as initramfs. For our boot path we need both.
echo "==> Adding initramfs essentials (/init, /dev/console)"

# /init script: mount devtmpfs (so /dev/console etc. exist before erlinit
# tries to use them) and exec the real init (erlinit at /sbin/init).
sudo tee "$WORK/rootfs/init" > /dev/null <<'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
sudo chmod +x "$WORK/rootfs/init"

# /dev/console — needed BEFORE devtmpfs is mounted for the redirect above.
sudo mkdir -p "$WORK/rootfs/dev"
sudo mknod -m 622 "$WORK/rootfs/dev/console" c 5 1

# 6: repack as cpio.gz (deterministic)
echo "==> Packing cpio.gz..."
( cd "$WORK/rootfs" && \
  sudo find . -mindepth 1 | LC_ALL=C sort | \
  sudo cpio --reproducible --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
        # Try the system's built host tree (Buildroot puts ubinize in host/sbin)
        for cand in \
            "${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 "WARNING: ubinize not found (install mtd-utils or build the system first), skipping UBI image generation" >&2
    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 "${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