#!/bin/bash
# Flash a TF-M image onto a Musca board through its on-board DAPLink USB mass
# storage, then let the DAPLink "Auto Reset" run it.
#
# Why not a USB power-cycle: the board is powered once and the target is reset
# in-band by the DAPLink after programming (DETAILS.TXT "Auto Reset: 1"). That
# keeps the serial console (ttyACM / ser2net) attached across the reset, so the
# boot is captured from the first byte, which the USB-power-cycle "musca" method
# loses. (The Musca DAPLink also used to wedge on a VBUS-only power-cycle by
# back-powering through the data lines; the LAA now cuts the data lines too.)
#
# Adapted from LAVA's "to: musca" deploy (mount / copy / sync / unmount). On the
# LAA there is a single board per appliance, so the DAPLink mass-storage node is
# auto-detected; set BOARD_SN to force a specific probe serial (ID_SERIAL_SHORT).
#
# Usage:
#   musca-flash-fw --flash <image>   # image: a .hex/.bin, or a tarball with one
set -eux

usage() { echo "Usage: $0 --flash <image>" >&2; exit 1; }

# The DAPLink mass-storage node, by serial if given, else the only one present.
find_msd() {
    if [ -n "${BOARD_SN:-}" ]; then
        echo "/dev/disk/by-id/usb-MBED_VFS_${BOARD_SN}-0:0"
        return 0
    fi
    ls /dev/disk/by-id/usb-MBED_VFS_*-0:0 2>/dev/null | head -n1
}

wait_path() {
    local p="$1" n="${2:-30}"
    while [ "$n" -gt 0 ]; do [ -e "$p" ] && return 0; n=$((n - 1)); sleep 1; done
    return 1
}

# Echo a single flashable .hex/.bin from the input (a raw image, or a tarball).
resolve_image() {
    local in="$1" d
    case "$in" in
        *.tar | *.tar.* | *.tgz | *.tbz2)
            d="$(mktemp -d)"
            tar xf "$in" -C "$d"
            find "$d" -type f \( -name '*.hex' -o -name '*.bin' \) | head -n1
            ;;
        *) printf '%s\n' "$in" ;;
    esac
}

# Program the image onto the DAPLink mass storage once. Returns 0 on success,
# 1 if the mount/copy fails or the DAPLink rejects the image (FAIL.TXT).
program_once() {
    local msd="$1" img="$2" mnt rc=0
    mnt="$(mktemp -d)"
    mount "$msd" "$mnt" || { rmdir "$mnt" 2>/dev/null || true; return 1; }
    cp "$img" "$mnt/" || rc=1
    sync          # DAPLink programs the flash during this flush and
    umount "$mnt" 2>/dev/null || rc=1 # finalises on unmount, then Auto-Resets
    rmdir "$mnt" 2>/dev/null || true
    [ "$rc" -eq 0 ] || return 1

    # The DAPLink remounts the mass storage after programming; a FAIL.TXT there
    # means the image was rejected. The remount lags (the /dev node reappears
    # before its media is ready, so a mount races with "no medium found"), so
    # retry the mount; a volume that never remounts is treated as success (a
    # good flash is simply mid reset-remount when we look).
    mnt="$(mktemp -d)"
    local mounted=
    for _ in $(seq 15); do
        if mount -o ro "$msd" "$mnt" 2>/dev/null; then mounted=1; break; fi
        sleep 2
    done
    if [ -n "$mounted" ]; then
        [ -e "$mnt/FAIL.TXT" ] && { cat "$mnt/FAIL.TXT" >&2; rc=1; }
        umount "$mnt" 2>/dev/null || true
    fi
    rmdir "$mnt" 2>/dev/null || true
    return "$rc"
}

flash() {
    local in="$1" img msd try=0
    [ -e "$in" ] || { echo "input '$in' not found" >&2; exit 1; }
    img="$(resolve_image "$in")"
    [ -n "$img" ] && [ -e "$img" ] || { echo "no flashable .hex/.bin in '$in'" >&2; exit 1; }

    # The Musca DAPLink occasionally rejects a program ("FAILED to download the
    # flash data contents"); retry the whole flash a few times.
    while :; do
        try=$((try + 1))
        msd="$(find_msd)"
        { [ -n "$msd" ] && wait_path "$msd" 30; } || { echo "DAPLink mass storage not found" >&2; exit 1; }
        program_once "$msd" "$img" && return 0
        [ "$try" -ge 3 ] && { echo "flash failed after ${try} attempts" >&2; exit 1; }
        echo "flash attempt ${try} rejected; retrying" >&2
        sleep 3
    done
}

case "${1:-}" in
    --flash) [ $# -ge 2 ] || usage; flash "$2" ;;
    *) usage ;;
esac