#!/usr/bin/env bash
# flash_macmesh_dualboot.sh — put both firmwares on a Heltec WiFi LoRa 32 V4,
# so one radio runs MeshCore or Meshtastic and the board picks at power-on.
#
# Everything it needs is downloaded from macmesh.xyzzy.computer and checked
# against the published SHA256.txt. Nothing else to fetch by hand, no repository
# to clone, no PlatformIO.
#
# Usage:
#   chmod +x flash_macmesh_dualboot.sh && ./flash_macmesh_dualboot.sh
#
# Options:
#   -p, --port PORT      Serial port. Default: the single /dev/cu.usbmodem*
#                        (macOS) or /dev/ttyACM*|/dev/ttyUSB* (Linux) found.
#   -y, --yes            Do not ask before erasing the board.
#   -u, --base-url URL   Where to fetch the images from. Default:
#                        https://macmesh.xyzzy.computer/downloads
#   -h, --help           This text.
#
# Requires: bash, curl, gunzip, shasum, python3 (for a throwaway esptool venv).
#
# What lands where, on a 16 MB board:
#
#   0x000000  base image     bootloader, partition table, Meshtastic application
#   0x020000  MacMesh Loader the selector that chooses between the two
#   0x0a0000  MeshCore       ota_0, UART GPIO47/48 at 57600 with BLE
#   0xb50000  Meshtastic FS  the serial settings the carrier needs
#
# The board keeps no identity of its own here: the Meshtastic filesystem has its
# private key stripped, and MeshCore makes a keypair on first boot from the radio
# chip's hardware RNG. Every board flashed with these images ends up distinct.

set -euo pipefail

BASE_URL="https://macmesh.xyzzy.computer/downloads"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CACHE_DIR="$SCRIPT_DIR/.macmesh-flash-cache"
VENV_DIR="$CACHE_DIR/venv"
PORT=""
ASSUME_YES=0

BASE_GZ="macmesh-dualboot-base-heltec-v4.bin.gz"
LOADER_BIN="macmesh-loader-heltec-v4.bin"
MESHCORE_BIN="macmesh-meshcore-dualboot-heltec-v4.bin"
MTFS_GZ="macmesh-meshtastic-fs-heltec-v4.bin.gz"

log() { echo "==> $*"; }
die() { echo "error: $*" >&2; exit 1; }

while [ $# -gt 0 ]; do
    case "$1" in
        -p|--port) PORT="${2:-}"; [ -n "$PORT" ] || die "--port needs a value"; shift 2 ;;
        -y|--yes) ASSUME_YES=1; shift ;;
        -u|--base-url) BASE_URL="${2:-}"; [ -n "$BASE_URL" ] || die "--base-url needs a value"; shift 2 ;;
        -h|--help) sed -n '2,32p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
        *) die "unknown argument: $1" ;;
    esac
done

for cmd in curl gunzip shasum python3; do
    command -v "$cmd" >/dev/null 2>&1 || die "$cmd is required but not installed"
done

mkdir -p "$CACHE_DIR"

# ---- download and check -----------------------------------------------------
# SHA256.txt covers every published file. A cached copy is reused only when its
# hash still matches, so a re-run after a new release fetches the new image
# rather than flashing a stale one.
log "Fetching checksums from $BASE_URL"
curl -fsSL -o "$CACHE_DIR/SHA256.txt" "$BASE_URL/SHA256.txt" \
    || die "could not fetch $BASE_URL/SHA256.txt"

expected_sha() {  # filename
    awk -v want="$1" '$2 == want { print $1; found = 1 } END { exit !found }' "$CACHE_DIR/SHA256.txt" \
        || die "$1 is not listed in SHA256.txt"
}

fetch() {  # filename
    local name="$1" want have
    want="$(expected_sha "$name")"
    if [ -f "$CACHE_DIR/$name" ]; then
        have="$(shasum -a 256 "$CACHE_DIR/$name" | awk '{print $1}')"
        [ "$have" = "$want" ] && { log "  $name (cached)"; return 0; }
    fi
    log "  $name"
    curl -fL --progress-bar -o "$CACHE_DIR/$name.partial" "$BASE_URL/$name" \
        || die "could not download $name"
    mv "$CACHE_DIR/$name.partial" "$CACHE_DIR/$name"
    have="$(shasum -a 256 "$CACHE_DIR/$name" | awk '{print $1}')"
    [ "$have" = "$want" ] || die "$name failed its checksum (expected $want, got $have)"
}

log "Downloading images"
fetch "$BASE_GZ"
fetch "$LOADER_BIN"
fetch "$MESHCORE_BIN"
fetch "$MTFS_GZ"

# The two large images ship compressed; esptool wants them plain.
gunzip -kfq "$CACHE_DIR/$BASE_GZ" || die "could not decompress $BASE_GZ"
gunzip -kfq "$CACHE_DIR/$MTFS_GZ" || die "could not decompress $MTFS_GZ"
BASE_IMG="$CACHE_DIR/${BASE_GZ%.gz}"
MTFS_IMG="$CACHE_DIR/${MTFS_GZ%.gz}"
LOADER_IMG="$CACHE_DIR/$LOADER_BIN"
MESHCORE_IMG="$CACHE_DIR/$MESHCORE_BIN"

# ---- toolchain --------------------------------------------------------------
if [ ! -x "$VENV_DIR/bin/python3" ] || ! "$VENV_DIR/bin/python3" -c 'import esptool' >/dev/null 2>&1; then
    log "Setting up esptool (one time, in $VENV_DIR)"
    python3 -m venv "$VENV_DIR" >/dev/null || die "could not create the virtualenv"
    "$VENV_DIR/bin/pip" install --quiet --upgrade pip >/dev/null 2>&1 || true
    "$VENV_DIR/bin/pip" install --quiet esptool || die "could not install esptool"
fi
VENV_PY="$VENV_DIR/bin/python3"

# ---- board ------------------------------------------------------------------
detect_port() {
    [ -n "$PORT" ] && { echo "$PORT"; return; }
    local c=() p
    for p in /dev/cu.usbmodem* /dev/ttyACM* /dev/ttyUSB*; do [ -e "$p" ] && c+=("$p"); done
    case "${#c[@]}" in
        0) die "no USB serial device found. Plug the radio in, and if it still
   does not appear, hold PRG while connecting it." ;;
        1) echo "${c[0]}" ;;
        *) die "more than one serial device (${c[*]}); choose with -p" ;;
    esac
}

# The V4 re-enumerates after every write, and the next call reliably hits a busy
# port, a missing device node, or the same board under a new name. Ten attempts
# five seconds apart covers it; fewer silently loses a whole write.
ESP_ATTEMPTS=10

esp() {
    local attempt out p
    for attempt in $(seq 1 "$ESP_ATTEMPTS"); do
        p="$(detect_port 2>/dev/null || true)"
        if [ -n "$p" ]; then
            out="$("$VENV_PY" -m esptool --chip esp32s3 --port "$p" --baud 921600 "$@" 2>&1)" \
                && { echo "$out"; return 0; }
        fi
        sleep 5
    done
    echo "${out:-no serial port appeared}" >&2
    return 1
}

write_verified() {  # offset image label
    local addr="$1" img="$2" lbl="$3" out
    log "$lbl -> $addr"
    esp write_flash --flash_mode dio --flash_freq 80m --flash_size 16MB "$addr" "$img" >/dev/null \
        || die "failed to write $lbl at $addr"
    out="$(esp verify_flash "$addr" "$img")" || die "could not verify $lbl at $addr"
    echo "$out" | grep -qi "verify OK" || die "$lbl at $addr does not match the image after writing"
    log "  verified"
}

log "Identifying board"
detect_port >/dev/null   # fails immediately, with advice, if nothing is attached
info="$(ESP_ATTEMPTS=2 esp flash_id)" \
    || die "found a serial port but no ESP32 answering on it. Hold PRG while
   plugging in USB, then run this again."
echo "$info" | sed -n 's/^\(Chip is\|MAC\|Detected flash size\).*/  &/p'
echo "$info" | grep -q "ESP32-S3" || die "this is not an ESP32-S3"
echo "$info" | grep -q "Detected flash size: 16MB" || die "the board reports less than 16MB of flash; both firmwares will not fit"
MAC="$(echo "$info" | sed -n 's/^MAC: //p' | tr -d ': ' | head -1)"

if [ "$ASSUME_YES" -ne 1 ]; then
    echo
    echo "About to ERASE and reflash the board on $(detect_port) (MAC $MAC)."
    echo "Any node identity already on it is destroyed."
    read -r -p "Continue? [y/N] " reply
    case "$reply" in [yY]|[yY][eE][sS]) ;; *) die "aborted" ;; esac
fi

log "Erasing flash"
esp erase_flash >/dev/null

# The base image spans 0x0-0xc50000 and carries its own stock MeshCore and a
# blank Meshtastic filesystem; writes 2 and 4 replace those regions afterwards.
# Verifying 0x0 against it would therefore fail by design, so only the three
# regions this script owns are checked.
log "1/4  base image (bootloader, partition table, Meshtastic application)"
esp write_flash --flash_mode dio --flash_freq 80m --flash_size 16MB 0x0 "$BASE_IMG" >/dev/null \
    || die "failed to write the base image at 0x0"
write_verified 0x20000  "$LOADER_IMG"   "2/4  MacMesh Loader"
write_verified 0xa0000  "$MESHCORE_IMG" "3/4  MeshCore"
write_verified 0xb50000 "$MTFS_IMG"     "4/4  Meshtastic settings"

log "Done: board $MAC carries both firmwares"
echo
echo "Reset the board. It boots the loader, waits about two seconds, then starts"
echo "whichever firmware ran last. Press PRG during those two seconds to switch."
echo "Either one talks to the Macintosh over the header UART on GPIO47/48."