#!/usr/bin/env bash
# Provision CM4 eMMC when reTerminal is connected in boot mode (eMMC disable jumper).
# On USB boot device (2b8e) detection: run rpiboot to expose eMMC, then wait for the user
# to choose Backup or Deploy in the portal — no auto-flash; action runs only after portal choice.
# Run this from udev or a systemd service. Requires: usbboot (rpiboot) built, golden image for deploy.

# Do NOT use set -e: rpiboot and intermediate commands may return non-zero without being fatal

# Redirect all output to log file (stdout + stderr) so nothing is lost to nohup buffering
LOG_FILE="${LOG_FILE:-/var/lib/cm4-provisioning/flash.log}"
exec >> "$LOG_FILE" 2>&1

# Load overrides from env file if present
[[ -f /opt/cm4-provisioning/env ]] && source /opt/cm4-provisioning/env

# Configuration - adjust paths and size for your setup
RPIBOOT_DIR="${RPIBOOT_DIR:-/opt/usbboot}"
GOLDEN_IMAGE="${GOLDEN_IMAGE:-/var/lib/cm4-provisioning/golden.img}"
# Expected eMMC size in bytes. reTerminal DM (CM4) has 32 GB eMMC (~31268536320 bytes).
EMMC_SIZE_BYTES="${EMMC_SIZE_BYTES:-$(( 32 * 1024 * 1024 * 1024 ))}"
LOG_TAG="cm4-flash"
STATUS_FILE="${STATUS_FILE:-/var/lib/cm4-provisioning/status.json}"
LOG_FILE="${LOG_FILE:-/var/lib/cm4-provisioning/flash.log}"
LOCK_FILE="${LOCK_FILE:-/var/lib/cm4-provisioning/flash.lock}"

log() { echo "[$(date -Iseconds)] [$LOG_TAG] $*"; logger -t "$LOG_TAG" "$*" 2>/dev/null || true; }

write_status() {
  local phase="$1" message="$2" progress="${3:-null}" error="${4:-}"
  local ts; ts=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
  message="${message//\"/\\\"}"; error="${error//\"/\\\"}"
  if [[ -n "$error" ]]; then
    printf '{"phase":"%s","message":"%s","progress":%s,"error":"%s","updated":"%s"}\n' \
      "$phase" "$message" "$progress" "$error" "$ts" > "$STATUS_FILE" 2>/dev/null || true
  else
    printf '{"phase":"%s","message":"%s","progress":%s,"updated":"%s"}\n' \
      "$phase" "$message" "$progress" "$ts" > "$STATUS_FILE" 2>/dev/null || true
  fi
}

# Prevent concurrent runs (udev fires multiple times during rpiboot re-enumeration)
exec 9>"$LOCK_FILE"
if ! flock -n 9; then
  echo "[$(date -Iseconds)] Another flash-emmc instance is already running; exiting."
  exit 0
fi
trap 'rm -f "$LOCK_FILE" "$CURRENT_DEVICE_FILE" "$DEVICE_SOURCE_FILE" 2>/dev/null; exec 9>&-' EXIT

# Optional: only run if this file exists (safety)
ENABLE_FILE="${ENABLE_FILE:-/etc/cm4-provisioning/enabled}"
if [[ -n "$ENABLE_FILE" && ! -f "$ENABLE_FILE" ]]; then
  log "Skipping: $ENABLE_FILE not present"
  write_status "idle" "Provisioning disabled (remove /etc/cm4-provisioning/enabled to enable)" "null" 2>/dev/null || true
  exit 0
fi

# When a device is detected we ask the user (dashboard): Backup or Deploy?
# These files are used to wait for the user's choice (written by dashboard, read by this script).
BACKUPS_DIR="${BACKUPS_DIR:-/var/lib/cm4-provisioning/backups}"
ACTION_REQUEST_FILE="${ACTION_REQUEST_FILE:-/var/lib/cm4-provisioning/action_request}"
CURRENT_DEVICE_FILE="${CURRENT_DEVICE_FILE:-/var/lib/cm4-provisioning/current_device}"
DEVICE_SOURCE_FILE="${DEVICE_SOURCE_FILE:-/var/lib/cm4-provisioning/device_source}"
WAIT_TIMEOUT="${WAIT_TIMEOUT:-600}"
# Golden image required only for Deploy; allow Backup without it
if [[ ! -f "$GOLDEN_IMAGE" ]]; then
  log "Golden image not found (Deploy will be unavailable): $GOLDEN_IMAGE"
fi

RPIBOOT_BIN="$RPIBOOT_DIR/rpiboot"
# Gadget dir for CM4 (64-bit); fallback to mass-storage-gadget
RPIBOOT_GADGET=""
for d in "$RPIBOOT_DIR/mass-storage-gadget64" "$RPIBOOT_DIR/mass-storage-gadget"; do
  if [[ -d "$d" ]]; then
    RPIBOOT_GADGET="$d"
    break
  fi
done
if [[ ! -x "$RPIBOOT_BIN" ]]; then
  log "rpiboot not found: $RPIBOOT_BIN (build usbboot and set RPIBOOT_DIR)"
  write_status "error" "rpiboot not installed" "null" "rpiboot not found. Run install-usbboot-on-host.sh on the host or build-and-deploy-usbboot-to-host.sh from your machine."
  exit 1
fi
if [[ -z "$RPIBOOT_GADGET" ]]; then
  log "rpiboot gadget dir not found under $RPIBOOT_DIR (need mass-storage-gadget64 or mass-storage-gadget)"
  write_status "error" "rpiboot gadget missing" "null" "Copy mass-storage-gadget(64) to $RPIBOOT_DIR"
  exit 1
fi
# rpiboot requires bootfiles.bin or one of bootcode*.bin in the gadget dir; empty dir causes "No 'bootcode' files found"
if [[ ! -f "$RPIBOOT_GADGET/bootfiles.bin" && ! -f "$RPIBOOT_GADGET/bootcode.bin" && ! -f "$RPIBOOT_GADGET/bootcode4.bin" && ! -f "$RPIBOOT_GADGET/bootcode5.bin" ]]; then
  log "rpiboot gadget dir has no boot files: $RPIBOOT_GADGET (reinstall usbboot)"
  write_status "error" "rpiboot gadget empty" "null" "No boot files in $RPIBOOT_GADGET. Reinstall usbboot: run install-usbboot-on-host.sh on the host or build-and-deploy-usbboot-to-host.sh from your machine."
  exit 1
fi

# Ensure status dir exists and start with running state
mkdir -p "$(dirname "$STATUS_FILE")" "$BACKUPS_DIR" 2>/dev/null || true
write_status "rpiboot" "Connecting to CM4 in boot mode…" "0"

# Block devices before rpiboot (so we can detect new one after)
before_devs=$(lsblk -nd -o NAME 2>/dev/null | sort)

log "Starting rpiboot to expose CM4 eMMC as mass storage..."
# Run rpiboot with 90s timeout so we don't hang if it doesn't exit cleanly when device switches to mass storage
rpiboot_exit=0
timeout 90 "$RPIBOOT_BIN" -d "$RPIBOOT_GADGET" || rpiboot_exit=$?
# timeout returns 124 if killed by timeout; 0 or other if rpiboot exited on its own
if [[ "$rpiboot_exit" -eq 124 ]]; then
  log "rpiboot timed out after 90s (device may have switched to mass storage)"
elif [[ "$rpiboot_exit" -ne 0 ]]; then
  log "rpiboot exited with code $rpiboot_exit"
  write_status "error" "rpiboot failed" "null" "rpiboot failed or no device connected. Check flash.log on host. Try unplug/replug USB."
  exit 1
fi

echo "[$(date -Iseconds)] rpiboot finished (exit=$rpiboot_exit); starting device scan"
log "rpiboot completed; waiting for block device..."
write_status "rpiboot" "rpiboot done, waiting for block device…" "10"

# rpiboot exits when device switches to mass storage; udev may need several seconds to create /dev/sdX
# Poll for new block device for up to 30s (device switch can be slow)
target_dev=""
for wait_sec in $(seq 2 2 10) $(seq 12 2 30); do
  sleep 2
  for dev in /dev/sd[a-z] /dev/sd[a-z][a-z]; do
    [[ -b "$dev" ]] || continue
    [[ "$dev" =~ [0-9]$ ]] && continue
    size=$(blockdev --getsize64 "$dev" 2>/dev/null || true)
    if [[ -n "$size" ]]; then
      if (( size >= EMMC_SIZE_BYTES * 95 / 100 && size <= EMMC_SIZE_BYTES * 105 / 100 )); then
        target_dev=$dev
        break 2
      fi
      if [[ -z "$target_dev" && "$before_devs" != *"${dev#/dev/}"* ]]; then
        target_dev=$dev
      fi
    fi
  done
  [[ -n "$target_dev" ]] && break
  log "Waiting for block device... ${wait_sec}s"
  write_status "rpiboot" "Waiting for eMMC block device… (${wait_sec}s)" "10"
done

log "Device scan complete. before_devs=[$before_devs] target_dev=[$target_dev]"
if [[ -z "$target_dev" ]]; then
  log "No suitable block device found after rpiboot (expected ~${EMMC_SIZE_BYTES} bytes)"
  write_status "error" "No eMMC device found" "null" "No suitable block device after rpiboot"
  exit 1
fi

# Ask user (dashboard): Backup or Deploy?
write_status "waiting_choice" "Device connected (USB boot mode). Choose Backup or Deploy in the dashboard." "null"
echo "usb" > "$DEVICE_SOURCE_FILE" 2>/dev/null || true
echo "$target_dev" > "$CURRENT_DEVICE_FILE" 2>/dev/null || true
log "Waiting for user choice (Backup or Deploy) in dashboard; timeout ${WAIT_TIMEOUT}s..."
for (( i = 0; i < WAIT_TIMEOUT; i += 2 )); do
  sleep 2
  if [[ -f "$ACTION_REQUEST_FILE" ]]; then
    action=$(cat "$ACTION_REQUEST_FILE" 2>/dev/null | tr -d '[:space:]' | tr '[:upper:]' '[:lower:]')
    rm -f "$ACTION_REQUEST_FILE" 2>/dev/null || true
    if [[ "$action" == "backup" ]]; then
      backup_name="backup-$(date +%Y%m%d-%H%M%S).img"
      backup_path="$BACKUPS_DIR/$backup_name"
      write_status "backup" "Creating backup…" "null"
      log "Backing up $target_dev to $backup_path..."
      if dd if="$target_dev" of="$backup_path" bs=4M status=progress conv=fsync; then
        log "Backup complete: $backup_path"
        write_status "done" "Backup complete: $backup_name" "100"
      else
        write_status "error" "Backup failed" "null" "dd failed"
      fi
    elif [[ "$action" == "deploy" ]]; then
      if [[ ! -f "$GOLDEN_IMAGE" ]]; then
        log "Golden image not found; cannot deploy."
        write_status "error" "Deploy unavailable" "null" "Golden image not found. Add golden.img to /var/lib/cm4-provisioning/ for deploy."
        exit 1
      fi
      write_status "flashing" "Writing golden image…" "null"
      log "Flashing $GOLDEN_IMAGE to $target_dev..."
      if dd if="$GOLDEN_IMAGE" of="$target_dev" bs=4M status=progress conv=fsync; then
        log "Flash complete. Remove eMMC disable jumper and power cycle the reTerminal."
        write_status "done" "Flash complete. Remove eMMC disable jumper and power cycle the reTerminal." "100"
      else
        write_status "error" "Flash failed" "null" "dd failed"
      fi
    else
      write_status "error" "Unknown action" "null" "action_request must be 'backup' or 'deploy'"
    fi
    exit 0
  fi
done
log "Timeout waiting for user choice"
write_status "idle" "Timeout waiting for choice. Connect the device again to retry." "null"
exit 0