#!/usr/bin/env bash
# flash.sh - Safe USB flash with explicit confirm.
# Refuses internal disks (NVMe + first SATA + boot devices).
set -euo pipefail

YES=0
ARGS=()
for a in "$@"; do
  case "$a" in
    -y|--yes) YES=1 ;;
    *) ARGS+=("$a") ;;
  esac
done
set -- "${ARGS[@]+${ARGS[@]}}"

DEV="${1:-}"
ISO="${2:-}"

if [[ -z "$DEV" || -z "$ISO" ]]; then
  echo "Usage: $0 [--yes] /dev/sdX path/to.iso"
  exit 1
fi

[[ -f "$ISO" ]] || { echo "ERR: iso not found: $ISO" >&2; exit 1; }
[[ -b "$DEV" ]] || { echo "ERR: not a block device: $DEV" >&2; exit 1; }

# Refuse obvious internal disks (nvme, first SATA, eMMC/SD, virtio).
case "$DEV" in
  /dev/nvme*|/dev/sda|/dev/mmcblk*|/dev/vd*)
    echo "ERR: refusing to write to likely-internal device: $DEV" >&2
    echo "     If you really mean it, use dd manually." >&2
    exit 2;;
esac

# Refuse if device is mounted as / or /boot.
# Use lsblk -no PKNAME to map partition → parent disk; sed 's/[0-9]*$//' breaks
# for /dev/nvme0n1pN and /dev/mmcblk0pN (would yield /dev/nvme0n / /dev/mmcblk0).
ROOT_PART=$(findmnt -n -o SOURCE /)
ROOT_PKNAME=$(lsblk -no PKNAME "$ROOT_PART" 2>/dev/null | head -1)
[[ -n "$ROOT_PKNAME" ]] && ROOT_SRC="/dev/$ROOT_PKNAME" || ROOT_SRC="$ROOT_PART"
[[ "$DEV" != "$ROOT_SRC" ]] || { echo "ERR: $DEV is the root device. Aborting." >&2; exit 2; }

SIZE=$(lsblk -bndo SIZE "$DEV" | head -1)
SIZE_GB=$((SIZE / 1024 / 1024 / 1024))
MODEL=$(lsblk -ndo MODEL,VENDOR "$DEV" | xargs)

cat <<EOF
=== About to wipe and flash ===
Device: $DEV
Size:   ${SIZE_GB} GiB
Model:  $MODEL
ISO:    $ISO

This DESTROYS all data on $DEV. Type 'yes' (lowercase) to continue.
EOF

if [[ "$YES" -eq 1 ]]; then
  echo "[--yes] skipping interactive confirm"
else
  read -r ANS
  [[ "$ANS" == "yes" ]] || { echo "Aborted."; exit 1; }
fi

echo "[*] Unmounting any partitions on $DEV..."
for p in $(lsblk -nro NAME "$DEV" | tail -n +2); do
  sudo umount "/dev/$p" 2>/dev/null || true
done

echo "[*] Flashing $ISO -> $DEV ..."
sudo dd if="$ISO" of="$DEV" bs=4M status=progress conv=fsync oflag=direct
sudo sync

# === Add a 3rd partition for install logs ===
# The flashed ISO is iso9660 + ESP (sda1, sda2) totaling ~759 MB on a much
# larger USB. Carve a 3rd MBR partition out of the remaining free space and
# mkfs.vfat as label S8N_LOGS. preseed early_command mounts this partition
# during install and writes logs to it; after a failed install, pull the USB,
# plug into onyx, and run scripts/read-usb-logs.sh /dev/sdX.
echo "[*] Re-reading partition table..."
sudo partprobe "$DEV" 2>&1 || true
sleep 2

# Find the highest end-sector of existing partitions; new partition starts after
ISO_END=$(sudo sfdisk -l -o End "$DEV" 2>/dev/null | awk '/^[ ]*[0-9]+/ {print $1}' | sort -n | tail -1)
ISO_END="${ISO_END:-0}"
[[ "$ISO_END" -gt 0 ]] || { echo "ERR: cannot read partition table on $DEV after dd" >&2; exit 3; }
LOG_START=$((ISO_END + 1))
# Round up to 1 MiB alignment (2048 sectors of 512B)
LOG_START=$(( (LOG_START + 2047) / 2048 * 2048 ))
TOTAL_SECTORS=$(sudo blockdev --getsz "$DEV")
LOG_SIZE=$((TOTAL_SECTORS - LOG_START - 2048))   # leave 1 MiB tail for safety
[[ "$LOG_SIZE" -gt 102400 ]] || { echo "ERR: not enough space for log partition (<50 MiB)" >&2; exit 3; }
LOG_SIZE_MB=$((LOG_SIZE / 2048))
echo "[*] Adding partition 3 at sector $LOG_START, size ${LOG_SIZE_MB} MiB ..."

# Use sfdisk to add a 3rd MBR entry without disturbing existing ones.
# `--no-reread` because the kernel already has the iso9660 partition open.
EXISTING_PT=$(sudo sfdisk -d "$DEV")
{
  echo "$EXISTING_PT"
  echo "${DEV}3 : start=$LOG_START, size=$LOG_SIZE, type=c"
} | sudo sfdisk --no-reread --no-tell-kernel "$DEV" || {
    # Fallback: write only the new entry by editing in place
    echo "[*] sfdisk batch failed; trying single-entry append"
    printf 'n\np\n3\n%s\n+%sK\nt\n3\nc\nw\n' "$LOG_START" "$((LOG_SIZE_MB * 1024))" \
      | sudo fdisk "$DEV" || true
  }
sudo partprobe "$DEV" 2>&1 || true
sleep 2

PART3="${DEV}3"
case "$DEV" in
  /dev/nvme*|/dev/mmcblk*) PART3="${DEV}p3" ;;
esac

if [[ -b "$PART3" ]]; then
  echo "[*] Formatting $PART3 as vfat (label S8N_LOGS)..."
  sudo mkfs.vfat -F 32 -n S8N_LOGS "$PART3"
  # Drop a README into the log partition so it's discoverable
  TMPDIR=$(mktemp -d)
  sudo mount "$PART3" "$TMPDIR"
  sudo tee "$TMPDIR/README.txt" >/dev/null <<EOF
S8N_LOGS partition
==================

This partition collects install logs during preseed/late_command on Debian
ISOs built by s8n/debian-s8ns-prefs-iso. If install fails on the target
machine, pull this USB, plug into another machine, and read /var/log/* +
/postinstall.log files here for diagnostics.

Companion script: scripts/read-usb-logs.sh /dev/sdX
EOF
  sudo umount "$TMPDIR"
  rmdir "$TMPDIR"
  echo "[OK] Log partition $PART3 ready (S8N_LOGS, vfat)"
else
  echo "[!] WARN: $PART3 didn't appear; log capture WILL NOT work on this USB"
fi

sudo sync
echo
echo "[OK] Done. Eject: sudo eject $DEV"
echo "[i] After install fail/success, mount $PART3 (label S8N_LOGS) to read logs"