#!/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.

set -e

# 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 CM4 often 8GB or 16GB). Used to identify the correct block device.
EMMC_SIZE_BYTES="${EMMC_SIZE_BYTES:-$(( 8 * 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}"

log() { echo "[$LOG_TAG] $*"; logger -t "$LOG_TAG" "$*"; echo "[$(date -Iseconds)] $*" >> "$LOG_FILE" 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
}

# 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

# 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..."
if ! "$RPIBOOT_BIN" -d "$RPIBOOT_GADGET"; then
  log "rpiboot failed or no device connected"
  write_status "error" "rpiboot failed" "null" "rpiboot failed or no device connected"
  exit 1
fi

# rpiboot exits when mass storage appears; give udev a moment to create /dev/sdX
sleep 3

# Find new block device (prefer one matching expected eMMC size)
target_dev=""
for dev in /dev/sd[a-z] /dev/sd[a-z][a-z]; do
  [[ -b "$dev" ]] || continue
  # Skip partitions
  [[ "$dev" =~ [0-9]$ ]] && continue
  size=$(blockdev --getsize64 "$dev" 2>/dev/null || true)
  if [[ -n "$size" ]]; then
    # Allow 5% tolerance on size
    if (( size >= EMMC_SIZE_BYTES * 95 / 100 && size <= EMMC_SIZE_BYTES * 105 / 100 )); then
      target_dev=$dev
      break
    fi
    # Otherwise take first new disk that appeared (fallback)
    if [[ -z "$target_dev" && "$before_devs" != *"${dev#/dev/}"* ]]; then
      target_dev=$dev
    fi
  fi
done

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 2>>"$LOG_FILE"; 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."
        rm -f "$CURRENT_DEVICE_FILE" "$DEVICE_SOURCE_FILE" 2>/dev/null || true
        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 2>>"$LOG_FILE"; 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
    rm -f "$CURRENT_DEVICE_FILE" "$DEVICE_SOURCE_FILE" 2>/dev/null || true
    exit 0
  fi
done
log "Timeout waiting for user choice"
write_status "idle" "Timeout waiting for choice. Connect the device again to retry." "null"
rm -f "$CURRENT_DEVICE_FILE" "$DEVICE_SOURCE_FILE" 2>/dev/null || true
exit 0