#!/usr/bin/env bash
# Revision: 2
# 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 (optional). If set, used to prefer among multiple new devices; if unset, any new block device after rpiboot is accepted (works for 8/16/32 GB CM4).
EMMC_SIZE_BYTES="${EMMC_SIZE_BYTES:-}"
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 (touch /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}"
# Optional: shrink backup with PiShrink (requires pishrink + parted, e2fsprogs). SHRINK_BACKUP=1 to enable.
SHRINK_BACKUP="${SHRINK_BACKUP:-0}"
# If SHRINK_BACKUP=1, optionally compress: PISHRINK_COMPRESS=gz or xz (uses parallel when available). Uncompressed .img can be dd'd directly for deploy.
PISHRINK_COMPRESS="${PISHRINK_COMPRESS:-}"
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, bootcode*.bin, or boot.img 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" && ! -f "$RPIBOOT_GADGET/boot.img" ]]; 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. On the host run: fix-gadget-bootcode-on-host.sh (or from your machine: ssh root@HOST 'bash -s' < scripts/fix-gadget-bootcode-on-host.sh). See docs troubleshooting."
  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 (gadget: $RPIBOOT_GADGET)..."
# Run rpiboot with 180s timeout so device has time to receive bootcode and switch to mass storage; -v for verbose
rpiboot_exit=0
timeout 180 "$RPIBOOT_BIN" -v -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 180s (device may have switched to mass storage)"
elif [[ "$rpiboot_exit" -ne 0 ]]; then
  log "rpiboot exited with code $rpiboot_exit"
  log "Common causes: (1) No device in USB boot mode — set eMMC disable jumper and use USB slave port. (2) Wrong USB port or cable. (3) Run on host: tail -50 /var/lib/cm4-provisioning/flash.log"
  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 new block device (any size — 8/16/32 GB CM4 supported)..."
write_status "rpiboot" "rpiboot done, waiting for block device…" "10"

# rpiboot exits when device switches to mass storage; udev may need many seconds to create /dev/sdX
# Dynamic detection: accept any NEW block device (not present before rpiboot). No fixed eMMC size required.
# When rpiboot timed out (124), device may still be switching — wait longer in that case
max_wait=60
[[ "$rpiboot_exit" -eq 124 ]] && max_wait=90
target_dev=""
new_devs=""
for wait_sec in $(seq 2 2 20) $(seq 22 2 $max_wait); do
  sleep 2
  new_devs=""
  for dev in /dev/sd[a-z] /dev/sd[a-z][a-z]; do
    [[ -b "$dev" ]] || continue
    [[ "$dev" =~ [0-9]$ ]] && continue
    # Only consider devices that appeared after rpiboot (the CM4 eMMC)
    if [[ "$before_devs" != *"${dev#/dev/}"* ]]; then
      size=$(blockdev --getsize64 "$dev" 2>/dev/null || true)
      [[ -n "$size" ]] && new_devs="$new_devs $dev:$size"
    fi
  done
  new_devs="${new_devs# }"
  if [[ -n "$new_devs" ]]; then
    # One new device: use it (dynamic — works for any eMMC size)
    # Multiple new devices: prefer one matching EMMC_SIZE_BYTES if set, else largest
    if [[ "$new_devs" != *" "* ]]; then
      target_dev="${new_devs%%:*}"
      break
    else
      best_dev="" best_size=0 best_delta=999999999999
      for entry in $new_devs; do
        dev="${entry%%:*}"
        size="${entry##*:}"
        if [[ -z "$EMMC_SIZE_BYTES" || "$EMMC_SIZE_BYTES" -eq 0 ]]; then
          # No size hint: take largest new device
          [[ "$size" -gt "$best_size" ]] && { best_dev="$dev"; best_size="$size"; }
        else
          delta=$(( size - EMMC_SIZE_BYTES )); [[ "$delta" -lt 0 ]] && delta=$(( -delta ))
          [[ "$delta" -lt "$best_delta" ]] && { best_dev="$dev"; best_delta="$delta"; }
        fi
      done
      [[ -n "$best_dev" ]] && target_dev="$best_dev" && break
    fi
  fi
  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 [[ -n "$target_dev" ]]; then
  detected_size=$(blockdev --getsize64 "$target_dev" 2>/dev/null || true)
  log "Using $target_dev (size=${detected_size:-?} bytes, $(( ${detected_size:-0} / 1024 / 1024 / 1024 )) GB)"
fi
if [[ -z "$target_dev" ]]; then
  log "No new block device found after rpiboot"
  log "Current block devices (for debugging):"
  lsblk -nd -o NAME,SIZE,TYPE 2>/dev/null | while read -r line; do log "  $line"; done
  for d in /dev/sd[a-z] /dev/sd[a-z][a-z]; do
    [[ -b "$d" ]] || continue
    [[ "$d" =~ [0-9]$ ]] && continue
    s=$(blockdev --getsize64 "$d" 2>/dev/null || true)
    log "  $d size=$s ($((${s:-0} / 1024 / 1024 / 1024)) GB)"
  done
  write_status "error" "No eMMC device found" "null" "No suitable block device after rpiboot. Check flash.log on host; unplug/replug and ensure eMMC disable jumper is set."
  exit 1
fi

# Ask user (dashboard): Backup or Deploy?
write_status "waiting_choice" "Device connected (USB boot mode). Choose Backup, Deploy, or Update EEPROM 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, Deploy, or Update EEPROM) 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
        final_name="$backup_name"
        shrink_requested=false
        [[ "$SHRINK_BACKUP" == "1" || "$SHRINK_BACKUP" == "true" ]] && shrink_requested=true
        [[ -f "$(dirname "$STATUS_FILE")/shrink_next_backup" ]] && shrink_requested=true
        rm -f "$(dirname "$STATUS_FILE")/shrink_next_backup" 2>/dev/null || true
        if [[ "$shrink_requested" == "true" ]]; then
          pishrink_cmd=""
          for p in /usr/local/bin/pishrink.sh /usr/local/bin/pishrink; do
            [[ -x "$p" ]] && pishrink_cmd="$p" && break
          done
          if [[ -n "$pishrink_cmd" ]]; then
            write_status "backup" "Shrinking backup image…" "null"
            log "Shrinking backup with PiShrink (timeout 30 min)..."
            pishrink_opts="-n"
            [[ "$PISHRINK_COMPRESS" == "gz" || "$PISHRINK_COMPRESS" == "gzip" ]] && pishrink_opts="$pishrink_opts -z -a"
            [[ "$PISHRINK_COMPRESS" == "xz" ]] && pishrink_opts="$pishrink_opts -Z -a"
            PISHRINK_TIMEOUT="${PISHRINK_TIMEOUT:-1800}"
            if timeout "$PISHRINK_TIMEOUT" $pishrink_cmd $pishrink_opts "$backup_path" 2>&1; then
              if [[ "$PISHRINK_COMPRESS" == "gz" || "$PISHRINK_COMPRESS" == "gzip" ]]; then
                final_name="${backup_name}.gz"
              elif [[ "$PISHRINK_COMPRESS" == "xz" ]]; then
                final_name="${backup_name}.xz"
              fi
              log "Shrunk backup: $BACKUPS_DIR/$final_name"
            else
              rc=$?
              if [[ "$rc" -eq 124 ]]; then
                log "PiShrink timed out after ${PISHRINK_TIMEOUT}s; keeping full backup $backup_path"
              else
                log "PiShrink failed (exit $rc); keeping full backup $backup_path"
              fi
            fi
          else
            log "SHRINK_BACKUP=1 but pishrink not found; keeping full backup. Install with scripts/install-pishrink-on-host.sh"
          fi
        fi
        log "Backup complete: $BACKUPS_DIR/$final_name"
        write_status "done" "Backup complete: $final_name" "100"
        ( sleep 90 && write_status "idle" "Waiting for Device in USB boot mode." "null" ) &
      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..."
      GOLDEN_RESOLVED="$(readlink -f "$GOLDEN_IMAGE" 2>/dev/null || echo "$GOLDEN_IMAGE")"
      if [[ "$GOLDEN_RESOLVED" == *.img.xz ]]; then
        decompress_cmd="xz -d -c"
      elif [[ "$GOLDEN_RESOLVED" == *.img.gz ]]; then
        decompress_cmd="gzip -c -d"
      else
        decompress_cmd="cat"
      fi
      if $decompress_cmd "$GOLDEN_IMAGE" 2>/dev/null | dd 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"
        # Auto-reset status to idle after 90s so dashboard does not stay on this message (dashboard also auto-clears after 60s when open)
        ( sleep 90 && write_status "idle" "Waiting for Device in USB boot mode." "null" ) &
      else
        write_status "error" "Flash failed" "null" "dd failed"
      fi
    elif [[ "$action" == "eeprom_update" ]]; then
      # Dashboard has written pieeprom.upd and pieeprom.sig to BASE_DIR; copy to eMMC boot partition
      PROV_DIR="$(dirname "$STATUS_FILE")"
      EEPROM_UPD="$PROV_DIR/pieeprom.upd"
      EEPROM_SIG="$PROV_DIR/pieeprom.sig"
      if [[ ! -f "$EEPROM_UPD" || ! -f "$EEPROM_SIG" ]]; then
        log "EEPROM update files not found: $EEPROM_UPD / $EEPROM_SIG"
        write_status "error" "EEPROM update failed" "null" "pieeprom.upd or pieeprom.sig not found in $PROV_DIR"
        exit 1
      fi
      boot_part=""
      for p in "${target_dev}1" "${target_dev}p1"; do
        if [[ -b "$p" ]]; then
          boot_part="$p"
          break
        fi
      done
      if [[ -z "$boot_part" ]]; then
        log "No boot partition found for $target_dev (tried ...1 and ...p1)"
        write_status "error" "EEPROM update failed" "null" "Boot partition not found"
        exit 1
      fi
      write_status "eeprom_update" "Writing EEPROM update to boot partition…" "null"
      log "Mounting $boot_part and copying EEPROM update..."
      mnt=$(mktemp -d)
      if mount "$boot_part" "$mnt" 2>/dev/null; then
        if cp "$EEPROM_UPD" "$mnt/pieeprom.upd" && cp "$EEPROM_SIG" "$mnt/pieeprom.sig"; then
          sync
          log "EEPROM update written. Remove eMMC disable jumper and power cycle to apply."
          write_status "done" "EEPROM update written to boot partition. Remove eMMC disable jumper and power cycle the reTerminal to apply." "100"
          rm -f "$EEPROM_UPD" "$EEPROM_SIG"
          ( sleep 90 && write_status "idle" "Waiting for Device in USB boot mode." "null" ) &
        else
          write_status "error" "EEPROM update failed" "null" "Failed to copy pieeprom.upd/sig"
        fi
        umount "$mnt" 2>/dev/null || true
      else
        write_status "error" "EEPROM update failed" "null" "Could not mount boot partition"
      fi
      rm -rf "$mnt"
    else
      write_status "error" "Unknown action" "null" "action_request must be 'backup', 'deploy', or 'eeprom_update'"
    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