#!/usr/bin/env bash
set -euo pipefail

IMAGE_PATH="${1:-viniciusos.img}"
PART_LABEL="VIN_DATA"

# --- 1. Image Check ---
if [[ ! -f "${IMAGE_PATH}" ]]; then
  echo "Error: OS image '${IMAGE_PATH}' not found." >&2
  exit 1
fi

# --- 2. Resolve Root Disk to Prevent Accidental Suicide ---
ROOT_DEV=$(findmnt -n -o SOURCE /)
ROOT_DISK=$(lsblk -no PKNAME "${ROOT_DEV}" 2>/dev/null || echo "${ROOT_DEV}" | sed -E 's/p?[0-9]+$//')
ROOT_DISK_NAME=$(basename "${ROOT_DISK}")

# --- 3. Scan for Target Removable / USB Disks ---
# Using -P (KEY="VALUE" pairs) to avoid space-splitting bugs on models like "Flash Disk"
mapfile -t USB_LIST < <(
  lsblk -d -P -o PATH,RM,TRAN,TYPE | \
  grep 'TYPE="disk"' | \
  grep -E 'RM="1"|TRAN="usb"' | \
  sed -n 's/.*PATH="\([^"]*\)".*/\1/p'
)

if [[ ${#USB_LIST[@]} -eq 0 ]]; then
  echo "No USB/removable drives detected." >&2
  exit 1
fi

echo "=== Detected Removable Drives ==="
for i in "${!USB_LIST[@]}"; do
  dev="${USB_LIST[$i]}"
  info=$(lsblk -d -n -o SIZE,MODEL "${dev}" | xargs)
  echo "  [$((i + 1))] ${dev} (${info})"
done
echo ""

# --- 4. Interactive Selection ---
if [[ ${#USB_LIST[@]} -eq 1 ]]; then
  TARGET_DEV="${USB_LIST[0]}"
  echo "Auto-selected sole USB drive: ${TARGET_DEV}"
else
  while true; do
    read -r -p "Select drive number (1-${#USB_LIST[@]}): " CHOICE
    if [[ "${CHOICE}" =~ ^[0-9]+$ ]] && (( CHOICE >= 1 && CHOICE <= ${#USB_LIST[@]} )); then
      TARGET_DEV="${USB_LIST[$((CHOICE - 1))]}"
      break
    fi
    echo "Invalid choice. Try again."
  done
fi

# Normalize: If user somehow passed a partition node, resolve to parent disk
if [[ "$(lsblk -no TYPE "${TARGET_DEV}")" == "part" ]]; then
  PARENT_NAME=$(lsblk -no PKNAME "${TARGET_DEV}")
  TARGET_DEV="/dev/${PARENT_NAME}"
  echo "Resolved partition to parent disk: ${TARGET_DEV}"
fi

# Hard lockout if target matches root system drive
if [[ "$(basename "${TARGET_DEV}")" == "${ROOT_DISK_NAME}" ]]; then
  echo "CRITICAL ERROR: Selected device ${TARGET_DEV} hosts the root filesystem! Aborting." >&2
  exit 1
fi

# --- 5. Unmount Any Active Mounts on Target ---
echo ""
echo "Unmounting active partitions on ${TARGET_DEV}..."
while IFS= read -r mnt; do
  if [[ -n "${mnt}" ]]; then
    echo "Unmounting ${mnt}..."
    sudo umount "${mnt}"
  fi
done < <(lsblk -no MOUNTPOINT "${TARGET_DEV}")

echo ""
lsblk "${TARGET_DEV}"
echo ""
read -r -p "Flash '${IMAGE_PATH}' and format FAT32 on ${TARGET_DEV}? ALL DATA WILL BE WIPED. [y/N]: " CONFIRM
if [[ "${CONFIRM}" != [yY] && "${CONFIRM}" != [yY][eE][sS] ]]; then
  echo "Aborted."
  exit 0
fi

# --- 6. Calculate Sector Alignment ---
IMG_SIZE_BYTES=$(stat -c%s "${IMAGE_PATH}")
IMG_SECTORS=$(( (IMG_SIZE_BYTES + 511) / 512 ))

# Align FAT32 start to next 1 MiB boundary (2048 sectors) after image
ALIGN_SECTORS=2048
if (( IMG_SECTORS <= ALIGN_SECTORS )); then
  START_SECTOR=${ALIGN_SECTORS}
else
  START_SECTOR=$(( ((IMG_SECTORS + ALIGN_SECTORS - 1) / ALIGN_SECTORS) * ALIGN_SECTORS ))
fi

echo "Writing '${IMAGE_PATH}' (${IMG_SIZE_BYTES} bytes, ${IMG_SECTORS} sectors)..."

# --- 7. Flash OS Image ---
sudo dd if="${IMAGE_PATH}" of="${TARGET_DEV}" bs=4M conv=fsync status=progress

# --- 8. Append MBR FAT32 Partition ---
echo "Creating FAT32 partition starting at LBA ${START_SECTOR}..."
echo "start=${START_SECTOR}, type=c" | sudo sfdisk --no-reread "${TARGET_DEV}"

# Force kernel partition table reload
sudo blockdev --rereadpt "${TARGET_DEV}" 2>/dev/null || sudo partprobe "${TARGET_DEV}"
sudo udevadm settle

if [[ "${TARGET_DEV}" =~ [0-9]$ ]]; then
  PART_DEV="${TARGET_DEV}p1"
else
  PART_DEV="${TARGET_DEV}1"
fi

# Wait for udev to actually create the device node (up to 5 seconds)
echo "Waiting for ${PART_DEV} to appear..."
TIMEOUT=10
while [[ ! -b "${PART_DEV}" && ${TIMEOUT} -gt 0 ]]; do
  sleep 0.5
  ((TIMEOUT--))
done

if [[ ! -b "${PART_DEV}" ]]; then
  echo "Error: Device node ${PART_DEV} failed to appear." >&2
  exit 1
fi

# --- 9. Format Partition ---
echo "Formatting ${PART_DEV} as FAT32..."
sudo mkfs.vfat -F 32 -n "${PART_LABEL}" "${PART_DEV}"

echo "Done! ${TARGET_DEV} is flashed and ready."