#!/bin/bash
# Erase the M5Stack Atom Echo over USB and write our own firmware (atom-echo.yaml).
# Everything goes through the ESP32's boot ROM, which no firmware can change, so the old
# firmware can't interfere. Plug the Echo straight into the computer (data cable, not a hub).
#
# Usage: bash flash.sh [MAC]
#   MAC: the Echo's Wi-Fi MAC (aa:bb:cc:dd:ee:ff). Without it, the script shows the chip's MAC and asks.
# Needs: ESPHome from pipx, and secrets.yaml next to this script with wifi_ssid and api_key.
# Asks for: sudo (to open the USB serial port) and the Wi-Fi password.
# Keeps: a copy of the old flash and the eFuse summary in ~/backups/atom-echo/ (mode 600).

set -uo pipefail
cd "$(dirname "$0")" || exit 1

EXPECTED_MAC=$(tr "A-F" "a-f" <<<"${1:-}")
VENV=~/.local/share/pipx/venvs/esphome/bin
ESPHOME="$VENV/esphome"
ESPTOOL="$VENV/esptool"
ESPEFUSE="$VENV/espefuse"
BAUD=460800
BACKUP_DIR=~/backups/atom-echo
STAMP=$(date +%Y%m%d-%H%M%S)

ok() { echo "🟢 $*"; }
info() { echo "🌱 $*"; }
fail() {
  echo "🔴 $*"
  echo
  echo "🔴 FAILED"
  exit 1
}

# esptool, retrying at 115200 if the faster speed fails (some Atom units don't like it)
et() {
  "$ESPTOOL" --port "$PORT" --baud "$BAUD" "$@" && return 0
  [ "$BAUD" = 115200 ] && return 1
  echo "🟠 failed at $BAUD baud, retrying at 115200"
  BAUD=115200
  "$ESPTOOL" --port "$PORT" --baud "$BAUD" "$@"
}

[ -x "$ESPHOME" ] || fail "ESPHome not installed (pipx install esphome)"
grep -q '^wifi_ssid:' secrets.yaml 2>/dev/null && grep -q '^api_key:' secrets.yaml \
  || fail "secrets.yaml needs wifi_ssid and api_key (see the README)"
SSID=$(sed -nE "s/^wifi_ssid: *['\"]?([^'\"]*)['\"]? *$/\1/p" secrets.yaml)

# 1. Find the Echo's USB serial port
mapfile -t PORTS < <(ls /dev/serial/by-id/* 2>/dev/null)
[ "${#PORTS[@]}" -eq 0 ] && fail "No USB serial device. Plug the Echo straight into the computer with a data cable (a hub may only power it)"
[ "${#PORTS[@]}" -gt 1 ] && fail "More than one USB serial device, unplug the others: ${PORTS[*]}"
PORT=$(readlink -f "${PORTS[0]}")
ok "USB serial port $PORT (${PORTS[0]##*/})"

# 2. Let this user open it (not in the dialout group)
if [ ! -w "$PORT" ]; then
  info "sudo is needed once to open $PORT"
  sudo setfacl -m "u:$USER:rw" "$PORT" || fail "Couldn't give $USER access to $PORT"
fi
ok "Port access"

# 3. Is it the right Echo? The MAC is burned into the chip's eFuses
MAC_OUT=$(et read-mac 2>&1) || { echo "$MAC_OUT"; fail "esptool couldn't talk to the chip"; }
MAC=$(grep -m1 -oiE 'MAC: *([0-9a-f]{2}:){5}[0-9a-f]{2}' <<<"$MAC_OUT" | awk '{print tolower($2)}')
if [ -n "$EXPECTED_MAC" ]; then
  [ "$MAC" = "$EXPECTED_MAC" ] || fail "MAC is '$MAC', expected $EXPECTED_MAC: not the Atom Echo we set up"
  ok "Chip MAC $MAC matches"
else
  read -rp "Chip MAC is $MAC. Erase and flash this device? [y/N] " YES
  [ "$YES" = y ] || [ "$YES" = Y ] || fail "Stopped: not confirmed"
  ok "Chip MAC $MAC confirmed"
fi

# 4. The old firmware mustn't have locked the chip (flash encryption or secure boot burned on)
mkdir -p "$BACKUP_DIR" && chmod 700 "$BACKUP_DIR"
EFUSE="$BACKUP_DIR/efuse-summary-$STAMP.txt"
"$ESPEFUSE" --port "$PORT" summary >"$EFUSE" 2>&1 || { cat "$EFUSE"; fail "espefuse summary failed"; }
chmod 600 "$EFUSE"
CRYPT=$(grep -m1 '^FLASH_CRYPT_CNT' "$EFUSE" | grep -oE '= [0-9]+' | tr -dc 0-9)
SB1=$(grep -m1 '^ABS_DONE_0' "$EFUSE" | grep -oE '= (True|False)')
SB2=$(grep -m1 '^ABS_DONE_1' "$EFUSE" | grep -oE '= (True|False)')
[ "$CRYPT" = 0 ] || fail "Flash encryption counter is '$CRYPT' (expected 0): the chip may be locked, stop here"
[ "$SB1" = "= False" ] && [ "$SB2" = "= False" ] || fail "Secure boot looks enabled ($SB1 / $SB2): stop here"
ok "No flash encryption, no secure boot (summary saved to $EFUSE)"

# 5. Keep a copy of the old firmware before erasing (may hold a Wi-Fi password, so mode 600)
OLD="$BACKUP_DIR/factory-flash-$STAMP.bin"
info "Copying the old 4 MB flash to $OLD (about 2 minutes)"
et read-flash 0 ALL "$OLD" || fail "Couldn't read the old flash"
chmod 600 "$OLD"
ok "Old flash saved, sha256 $(sha256sum "$OLD" | cut -c1-16)…"

# 6. Wi-Fi password into secrets.yaml (git-ignored, mode 600)
read -rsp "Wi-Fi password for $SSID: " PW
echo
[ -n "$PW" ] || fail "Empty password"
umask 077
ESC=${PW//\'/\'\'}
{ grep -v '^wifi_password:' secrets.yaml; printf "wifi_password: '%s'\n" "$ESC"; } >secrets.yaml.new && mv secrets.yaml.new secrets.yaml
unset PW ESC
ok "Wi-Fi password saved in secrets.yaml"

# 7. Build with it (memory capped, so a big build can't freeze a computer with little swap)
info "Building the firmware"
systemd-run --user --scope -q -p MemoryMax=6G "$ESPHOME" compile atom-echo.yaml >"$BACKUP_DIR/compile-$STAMP.log" 2>&1 \
  || { tail -30 "$BACKUP_DIR/compile-$STAMP.log"; fail "Build failed (log: $BACKUP_DIR/compile-$STAMP.log)"; }
IMG=$(find .esphome/build -name firmware.factory.bin -newer secrets.yaml | head -1)
[ -f "$IMG" ] || fail "No firmware.factory.bin from the build"
ok "Built $IMG, sha256 $(sha256sum "$IMG" | cut -c1-16)…"

# 8. Erase everything, write ours, and have the boot ROM check it
info "Erasing the whole flash chip"
et erase-flash || fail "Erase failed"
ok "Flash erased"
info "Writing our firmware"
et write-flash 0x0 "$IMG" || fail "Write failed"
et verify-flash 0x0 "$IMG" || fail "Verify failed: flash doesn't match our build"
ok "Flash verified against our build"

echo
ok "SUCCESS: the Echo now runs our firmware ($(grep -m1 -E '^    version:' atom-echo.yaml | tr -d ' "')). Next: the router rules and adding it to Home Assistant (README)"