#!/usr/bin/env bash
# flash_heltec_v4.sh — install stock Meshtastic firmware onto an attached
# Heltec WiFi LoRa 32 V4 and apply the Serial Module configuration for the
# MacMesh carrier's header UART (Rev7+ THVD1451 wiring: GPIO47 = TX, GPIO48 = RX).
#
# Background: docs/software-meshtastic.md and AGENTS.md. No custom firmware
# fork is needed for Meshtastic — this script just flashes the official
# heltec-v4 (OLED) release and configures the stock Serial Module.
#
# Usage:
#   scripts/flash_heltec_v4.sh [options]
#
# Options:
#   -p, --port PORT     Serial port to use. Default: auto-detect a single
#                        /dev/cu.usbmodem* (macOS) or /dev/ttyACM*|/dev/ttyUSB*
#                        (Linux) device. Errors out if zero or more than one
#                        candidate is found — pass -p to disambiguate.
#   -v, --version TAG    Meshtastic firmware release tag to flash.
#                        Default: v2.7.26.54e0d8d (pinned in
#                        docs/software-meshtastic.md; matches this project's
#                        verified serial-module behavior).
#   -b, --board NAME     PlatformIO target name. Default: heltec-v4.
#                        Do NOT use heltec-v4-tft — it reassigns GPIO44 to
#                        touch reset and disables the Serial Module in
#                        color-display mode (see docs/software-meshtastic.md).
#   --skip-flash         Skip erase/flash; only apply the Serial Module config.
#   --skip-config        Flash only; skip the Serial Module configuration step.
#   --skip-time          Don't push this computer's clock to the node. The
#                        Heltec has no GPS and never pairs with a phone app
#                        here, so it has no other time source — by default
#                        every run sets node time from this machine's clock
#                        (accurate; the classic Mac's clock is irrelevant, it
#                        never reaches the node). Note the ESP32-S3 has no
#                        battery-backed RTC, so this is lost on power loss and
#                        needs resetting again once the node is unplugged.
#   --macmesh            Switch the mode/baud/timeout to the TEXTMSG profile
#                        documented in docs/software-meshtastic.md for the
#                        classic-Mac MacMeshMeshtastic.app (--mode TEXTMSG
#                        --baud BAUD_9600 --timeout 0) instead of the default
#                        PROTO/API profile below. Does not touch --rxd/--txd —
#                        those are a hardware fact of this carrier revision,
#                        not a protocol choice.
#   --mode MODE          serial.mode. Default: PROTO (the protobuf/API serial
#                        interface used by the meshtastic CLI/Python API).
#                        Pass --macmesh (or --mode TEXTMSG --baud BAUD_9600)
#                        for MacMesh's classic Mac TEXTMSG app instead — pick
#                        one, not both.
#   --baud BAUD          serial.baud enum. Default: BAUD_38400.
#   --rxd PIN            serial.rxd. Default: 48 (Heltec RX <- Mac TX; Rev7+
#                        THVD1451 carrier wiring — GPIO47/48, not the older
#                        Rev4-Rev6 MAX3490E carrier's GPIO43/44).
#   --txd PIN            serial.txd. Default: 47 (Heltec TX -> Mac RX; Rev7+).
#   --timeout N          serial.timeout. Default: 0.
#   --echo BOOL          serial.echo (true|false). Default: false.
#   --region REGION      lora.region. Default: US. Required before the radio
#                        will transmit at all — a freshly flashed node's
#                        region is UNSET and stays receive-only/silent until
#                        this is set. Other values: EU_433, EU_868, CN, JP,
#                        ANZ, KR, TW, RU, IN, NZ_865, TH, UA_433, UA_868, ...
#                        (see `meshtastic --get lora.region` help for the
#                        full enum this build ships).
#   --hop-limit N        lora.hop_limit. Default: 7 (the firmware max; the
#                        Meshtastic default is 3). Each hop is a full
#                        retransmission, so higher trades more airtime/
#                        collision risk for reach across a wider mesh — 7 is
#                        the right call for a node that joins a shared
#                        community mesh rather than talking point-to-point.
#   -y, --yes            Don't prompt for confirmation before erasing flash.
#   -h, --help           Show this help and exit.
#
# The pin directions above are from the Heltec's point of view and match the
# Rev7+ MacMesh carrier wiring in AGENTS.md (THVD1451, GPIO47/48). This is
# NOT the same pinout as the older Rev4-Rev6 MAX3490E carrier (GPIO43/44) —
# Rev7+ deliberately moved off GPIO43/44 because those are also the ESP32-S3
# ROM bootloader's UART0 pins, so a Rev4-Rev6 board sees one garbled line of
# boot banner on every reset; a Rev7+ board on GPIO47/48 does not.
#
# Requires: bash, curl, unzip, jq, shasum, python3 (used only to create a
# throwaway virtualenv with esptool + the meshtastic CLI on first run).

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CACHE_DIR="$SCRIPT_DIR/.heltec-flash-cache"
VENV_DIR="$CACHE_DIR/venv"
DOWNLOAD_DIR="$CACHE_DIR/downloads"
EXTRACT_ROOT="$CACHE_DIR/extracted"

FW_VERSION="v2.7.26.54e0d8d"
FW_VERSION_PINNED_SHA256="339650ced8b8c859751f04be69c6af845fea96baca1327518f199afd2595a9d4"
BOARD="heltec-v4"
PORT=""
DO_FLASH=1
DO_CONFIG=1
DO_SET_TIME=1
ASSUME_YES=0

SERIAL_MODE="PROTO"
SERIAL_BAUD="BAUD_38400"
SERIAL_RXD="48"
SERIAL_TXD="47"
SERIAL_TIMEOUT="0"
SERIAL_ECHO="false"
REGION="US"
HOP_LIMIT="7"

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

usage() { sed -n '2,79p' "$0" | sed 's/^# \{0,1\}//'; }

while [ $# -gt 0 ]; do
    case "$1" in
        -p|--port) PORT="$2"; shift 2 ;;
        -v|--version) FW_VERSION="$2"; shift 2 ;;
        -b|--board) BOARD="$2"; shift 2 ;;
        --skip-flash) DO_FLASH=0; shift ;;
        --skip-config) DO_CONFIG=0; shift ;;
        --skip-time) DO_SET_TIME=0; shift ;;
        --macmesh) SERIAL_MODE="TEXTMSG"; SERIAL_BAUD="BAUD_9600"; SERIAL_TIMEOUT="0"; shift ;;
        --mode) SERIAL_MODE="$2"; shift 2 ;;
        --baud) SERIAL_BAUD="$2"; shift 2 ;;
        --rxd) SERIAL_RXD="$2"; shift 2 ;;
        --txd) SERIAL_TXD="$2"; shift 2 ;;
        --timeout) SERIAL_TIMEOUT="$2"; shift 2 ;;
        --echo) SERIAL_ECHO="$2"; shift 2 ;;
        --region) REGION="$2"; shift 2 ;;
        --hop-limit) HOP_LIMIT="$2"; shift 2 ;;
        -y|--yes) ASSUME_YES=1; shift ;;
        -h|--help) usage; exit 0 ;;
        *) die "unknown argument: $1 (see --help)" ;;
    esac
done

for cmd in curl unzip jq shasum python3; do
    command -v "$cmd" >/dev/null 2>&1 || die "required tool '$cmd' not found on PATH"
done

# ---------------------------------------------------------------------------
# Port detection. The Heltec's native USB-CDC port name changes between its
# ROM-bootloader identity and its running-app identity (observed on macOS as
# /dev/cu.usbmodem<MAC-suffix> vs /dev/cu.usbmodem<number>) — always re-glob
# right before touching the port rather than caching a stale path.
# ---------------------------------------------------------------------------
detect_port() {
    if [ -n "$PORT" ]; then
        echo "$PORT"
        return
    fi
    local candidates=()
    for p in /dev/cu.usbmodem* /dev/cu.SLAB_USBtoUART* /dev/cu.wchusbserial* /dev/ttyACM* /dev/ttyUSB*; do
        [ -e "$p" ] && candidates+=("$p")
    done
    case "${#candidates[@]}" in
        0) die "no USB serial device found. Is the Heltec V4 plugged in?" ;;
        1) echo "${candidates[0]}" ;;
        *) die "multiple USB serial devices found (${candidates[*]}); pass -p/--port to pick one" ;;
    esac
}

# Run an esptool/meshtastic command, re-resolving the port fresh each attempt
# and retrying on the transient "device not configured"/port-busy errors that
# show up around the Heltec's USB re-enumeration. The first three arguments
# are always the fixed "<python3> -m <module>" prefix (see ESPTOOL/MESHTASTIC
# below) — --port is inserted right after it, since esptool's argparse (unlike
# meshtastic's) requires --port before the subcommand, not after it.
with_port_retry() {
    local py modflag mod
    py="$1"; modflag="$2"; mod="$3"; shift 3
    local attempt out rc
    for attempt in 1 2 3; do
        local p; p="$(detect_port)"
        set +e
        out="$("$py" "$modflag" "$mod" --port "$p" "$@" 2>&1)"
        rc=$?
        set -e
        if [ $rc -eq 0 ]; then
            echo "$out"
            return 0
        fi
        if echo "$out" | grep -qi "device not configured\|port is busy\|no such file or directory\|not found"; then
            sleep 2
            continue
        fi
        echo "$out" >&2
        return $rc
    done
    echo "$out" >&2
    return 1
}

# ---------------------------------------------------------------------------
# Toolchain: a throwaway venv with esptool + the meshtastic CLI, cached under
# scripts/.heltec-flash-cache (gitignored) so repeat runs don't reinstall.
# ---------------------------------------------------------------------------
setup_venv() {
    if ! "$VENV_DIR/bin/python3" -c 'import esptool, meshtastic' >/dev/null 2>&1; then
        log "Setting up esptool/meshtastic CLI virtualenv at $VENV_DIR"
        python3 -m venv "$VENV_DIR"
        "$VENV_DIR/bin/pip" install --quiet --upgrade pip
        "$VENV_DIR/bin/pip" install --quiet esptool meshtastic
    fi
}

# Invoked as "$VENV_DIR/bin/python3" -m <module> rather than the pip-generated
# esptool.py/meshtastic wrapper scripts: this repo's own path contains a space
# ("Macintosh Meshtastic"), and a shebang line can't carry a space in its
# interpreter path — the OS truncates it there and the wrapper script fails to
# exec. Calling the module directly through python3 sidesteps that entirely
# (the same reason Meshtastic's own device-install.sh prefers `python -m esptool`).
ESPTOOL=()
MESHTASTIC=()

flash_firmware() {
    log "Resolving Meshtastic firmware release $FW_VERSION for board $BOARD"
    local api_url="https://api.github.com/repos/meshtastic/firmware/releases/tags/$FW_VERSION"
    local zip_url
    zip_url="$(curl -sL "$api_url" | jq -r '.assets[] | select(.name == "firmware-esp32s3-'"${FW_VERSION#v}"'.zip") | .browser_download_url')"
    [ -n "$zip_url" ] || die "could not find firmware-esp32s3-${FW_VERSION#v}.zip in release $FW_VERSION"

    mkdir -p "$DOWNLOAD_DIR"
    local zip_path="$DOWNLOAD_DIR/firmware-esp32s3-${FW_VERSION#v}.zip"
    if [ ! -f "$zip_path" ]; then
        log "Downloading $(basename "$zip_path") (~160MB, from github.com/meshtastic/firmware)"
        curl -L --progress-bar -o "$zip_path.partial" "$zip_url"
        mv "$zip_path.partial" "$zip_path"
    else
        log "Using cached download: $zip_path"
    fi

    local actual_sha256
    actual_sha256="$(shasum -a 256 "$zip_path" | cut -d' ' -f1)"
    if [ "$FW_VERSION" = "v2.7.26.54e0d8d" ]; then
        [ "$actual_sha256" = "$FW_VERSION_PINNED_SHA256" ] || die "checksum mismatch on $zip_path (expected $FW_VERSION_PINNED_SHA256, got $actual_sha256) — re-download before flashing"
    else
        log "Note: no pinned checksum for $FW_VERSION; downloaded sha256 is $actual_sha256"
    fi

    local extract_dir="$EXTRACT_ROOT/${FW_VERSION#v}"
    if [ ! -d "$extract_dir" ]; then
        log "Extracting firmware zip"
        mkdir -p "$extract_dir"
        unzip -q "$zip_path" -d "$extract_dir"
    fi

    local prog="firmware-${BOARD}-${FW_VERSION#v}"
    local factory_bin="$extract_dir/${prog}.factory.bin"
    local meta_json="$extract_dir/${prog}.mt.json"
    [ -f "$factory_bin" ] || die "factory image not found for board '$BOARD': $factory_bin"
    [ -f "$meta_json" ] || die "metadata json not found: $meta_json"

    local mcu ota_offset spiffs_offset
    mcu="$(jq -r '.mcu' "$meta_json")"
    ota_offset="$(jq -r '.part[] | select(.subtype == "ota_1") | .offset' "$meta_json")"
    spiffs_offset="$(jq -r '.part[] | select(.subtype == "spiffs") | .offset' "$meta_json")"
    local ota_bin="$extract_dir/mt-${mcu}-ota.bin"
    local spiffs_bin="$extract_dir/littlefs-${BOARD}-${FW_VERSION#v}.bin"
    [ -f "$ota_bin" ] || die "OTA image not found: $ota_bin"
    [ -f "$spiffs_bin" ] || die "LittleFS image not found: $spiffs_bin"

    local port; port="$(detect_port)"
    log "Identifying chip on $port before touching flash"
    with_port_retry "${ESPTOOL[@]}" chip_id

    if [ "$ASSUME_YES" -ne 1 ]; then
        echo
        echo "About to ERASE and reflash the device on $port with:"
        echo "  firmware : $prog ($FW_VERSION)"
        echo "  board    : $BOARD"
        read -r -p "Continue? [y/N] " reply
        case "$reply" in
            [yY]|[yY][eE][sS]) ;;
            *) die "aborted" ;;
        esac
    fi

    log "Erasing flash"
    with_port_retry "${ESPTOOL[@]}" erase_flash

    log "Writing factory image at 0x00"
    with_port_retry "${ESPTOOL[@]}" write_flash 0x00 "$factory_bin"

    log "Writing OTA slot at $ota_offset"
    with_port_retry "${ESPTOOL[@]}" write_flash "$ota_offset" "$ota_bin"

    log "Writing LittleFS at $spiffs_offset"
    with_port_retry "${ESPTOOL[@]}" write_flash "$spiffs_offset" "$spiffs_bin"

    log "Flash complete; waiting for the device to boot"
    sleep 5
}

configure_serial() {
    log "Applying Serial Module configuration and LoRa settings"
    log "  mode=$SERIAL_MODE baud=$SERIAL_BAUD rxd=$SERIAL_RXD txd=$SERIAL_TXD timeout=$SERIAL_TIMEOUT echo=$SERIAL_ECHO region=$REGION hop_limit=$HOP_LIMIT"
    with_port_retry "${MESHTASTIC[@]}" \
        --set serial.enabled true \
        --set serial.mode "$SERIAL_MODE" \
        --set serial.rxd "$SERIAL_RXD" \
        --set serial.txd "$SERIAL_TXD" \
        --set serial.baud "$SERIAL_BAUD" \
        --set serial.timeout "$SERIAL_TIMEOUT" \
        --set serial.echo "$SERIAL_ECHO" \
        --set serial.override_console_serial_port false \
        --set lora.region "$REGION" \
        --set lora.hop_limit "$HOP_LIMIT"

    sleep 3
    log "Reading back Serial Module + LoRa configuration"
    with_port_retry "${MESHTASTIC[@]}" --get serial
    with_port_retry "${MESHTASTIC[@]}" --get lora.region
    with_port_retry "${MESHTASTIC[@]}" --get lora.hop_limit
}

# The Heltec has no GPS and never pairs with a phone app (its only client is
# the classic Mac over TEXTMSG/PROTO), so it has no other source of correct
# time — the classic Mac's own clock is irrelevant here; it never reaches the
# node. This pushes THIS computer's clock instead. --set-time is a standalone
# admin command (mutually exclusive with --set in meshtastic's own CLI), so
# it runs as its own call, after the config transaction above commits.
sync_time() {
    log "Setting node time from this computer's clock ($(date))"
    with_port_retry "${MESHTASTIC[@]}" --set-time
}

main() {
    setup_venv
    ESPTOOL=("$VENV_DIR/bin/python3" -m esptool)
    MESHTASTIC=("$VENV_DIR/bin/python3" -m meshtastic)

    [ "$DO_FLASH" -eq 1 ] && flash_firmware
    [ "$DO_CONFIG" -eq 1 ] && configure_serial
    [ "$DO_SET_TIME" -eq 1 ] && sync_time

    log "Final device info"
    with_port_retry "${MESHTASTIC[@]}" --info | grep -i -A1 "firmwareVersion\|hwModel\|pioEnv" || true

    log "Done."
}

main