#!/bin/bash
VERSION=$(dpkg-query -W -f='${Version}' mw-pxeboot 2>/dev/null || echo "unknown")

usage() {
  cat <<'EOF'
Usage: mw-pxeboot [OPTIONS]

Set the next boot to PXE and reboot (default action). Always asks for
confirmation first unless --force is given -- this reboots the host
into network install, which will wipe and reinstall it.

Options:
  -h, --help          Show this help and exit
  -v, --version       Show version and exit
  --purge-non-pxe     Additionally delete all UEFI NVRAM boot entries
                      except the detected PXE/network boot entry before
                      rebooting. The confirmation prompt lists exactly
                      what will be deleted. Does not reboot at all if
                      any entry failed to delete.
  --force             Skip the confirmation prompt

Examples:
  mw-pxeboot                  Confirm, then set next boot to PXE and reboot
  mw-pxeboot --purge-non-pxe  Confirm, then also remove stale non-PXE
                               UEFI NVRAM boot entries before rebooting
EOF
}

FORCE=0
ACTION="boot-pxe"

for arg in "$@"; do
  case "$arg" in
    -h|--help)
      usage
      exit 0
      ;;
    -v|--version|version)
      echo "mw-pxeboot $VERSION"
      exit 0
      ;;
    --purge-non-pxe)
      ACTION="purge-non-pxe"
      ;;
    --force)
      FORCE=1
      ;;
    *)
      echo "Error: unknown option '$arg'" >&2
      usage >&2
      exit 1
      ;;
  esac
done

[ "$EUID" -ne 0 ] && { echo "Error: Must run as root"; exit 1; }

confirm() {
  local prompt="$1"
  if [ "$FORCE" -eq 1 ]; then
    return 0
  fi
  if [ ! -t 0 ]; then
    echo "Error: confirmation required but no interactive terminal available. Re-run with --force." >&2
    exit 1
  fi
  read -r -p "$prompt [y/N] " reply
  case "$reply" in
    [yY]|[yY][eE][sS]) return 0 ;;
    *) echo "Aborted."; exit 1 ;;
  esac
}

find_pxe_bootnums() {
  efibootmgr | grep -E '^Boot[0-9A-F]{4}' | grep -i "ip4" | sed 's/^Boot\([0-9A-F]*\).*/\1/'
  efibootmgr | grep -E '^Boot[0-9A-F]{4}' | grep -i "ip6" | sed 's/^Boot\([0-9A-F]*\).*/\1/'
  efibootmgr | grep -E '^Boot[0-9A-F]{4}' | grep -iE "network|pxe|gbe|ethernet" | grep -viE "ip4|ip6" | sed 's/^Boot\([0-9A-F]*\).*/\1/'
}

find_pxe_bootnum() {
  find_pxe_bootnums | head -1
}

do_boot_pxe() {
  echo "Setting next boot to PXE..."
  if ipmitool chassis bootdev pxe options=efiboot >/dev/null 2>&1; then
    echo "Next boot: PXE"
    sleep 2
    reboot -f
  elif command -v efibootmgr >/dev/null 2>&1; then
    local pxe
    pxe=$(find_pxe_bootnum)
    if [ -n "$pxe" ]; then
      efibootmgr --bootnext "$pxe" && echo "Next boot: PXE (Boot$pxe)" && sleep 2 && reboot -f
    else
      echo "Error: No PXE entry found"
      exit 1
    fi
  else
    echo "Error: No supported method (ipmitool/efibootmgr)"
    exit 1
  fi
}

case "$ACTION" in
  purge-non-pxe)
    command -v efibootmgr >/dev/null 2>&1 || { echo "Error: efibootmgr not available"; exit 1; }

    # Detect PXE entries: IPv4 first, then IPv6, then any other network/pxe-named entry.
    mapfile -t PXE_LIST < <(find_pxe_bootnums | awk '!seen[$0]++')
    if [ "${#PXE_LIST[@]}" -eq 0 ]; then
      echo "Error: No PXE entry found -- refusing to purge (would leave no PXE fallback)"
      exit 1
    fi
    echo "PXE boot entries: ${PXE_LIST[*]} (will be kept)"

    # Reorder BootOrder so the PXE entries come first, IPv4 before IPv6.
    ORDER_LINE=$(efibootmgr | sed -n 's/^BootOrder: //p')
    mapfile -t -d ',' ORDER_NUMS < <(printf '%s' "$ORDER_LINE")
    NEW_ORDER=("${PXE_LIST[@]}")
    for o in "${ORDER_NUMS[@]}"; do
      keep=1
      for p in "${PXE_LIST[@]}"; do
        if [ "$o" = "$p" ]; then keep=0; break; fi
      done
      [ "$keep" -eq 1 ] && NEW_ORDER+=("$o")
    done
    NEW_ORDER_STR=$(IFS=,; echo "${NEW_ORDER[*]}")
    if [ "$NEW_ORDER_STR" != "$ORDER_LINE" ]; then
      echo "Reordering BootOrder: $NEW_ORDER_STR"
      efibootmgr -o "$NEW_ORDER_STR"
    else
      echo "BootOrder already PXE-first: $ORDER_LINE"
    fi

    # Attempt to delete every non-PXE NVRAM boot entry.
    OTHERS=()
    while read -r b; do
      skip=0
      for p in "${PXE_LIST[@]}"; do
        if [ "$b" = "$p" ]; then skip=1; break; fi
      done
      [ "$skip" -eq 0 ] && OTHERS+=("$b")
    done < <(efibootmgr | grep -E '^Boot[0-9A-F]{4}' | sed 's/^Boot\([0-9A-F]*\).*/\1/')
    if [ "${#OTHERS[@]}" -eq 0 ]; then
      echo "No non-PXE NVRAM boot entries found."
      exit 0
    fi
    echo "The following non-PXE NVRAM boot entries will be deleted, then the host will reboot to PXE:"
    for b in "${OTHERS[@]}"; do
      efibootmgr | grep "^Boot${b}"
    done
    confirm "Delete these ${#OTHERS[@]} non-PXE NVRAM boot entries and reboot to PXE?"
    FAILED=0
    for b in "${OTHERS[@]}"; do
      if efibootmgr -b "$b" -B >/dev/null 2>&1; then
        echo "Deleted Boot${b}"
      else
        echo "Failed to delete Boot${b} (still referenced in BootOrder, or firmware refused the delete)" >&2
        FAILED=1
      fi
    done
    if [ "$FAILED" -eq 1 ]; then
      echo "Some entries could not be deleted via the standard efibootmgr method." >&2
      echo "If any entry remains (e.g. the firmware refuses to delete a Boot#### entry still" >&2
      echo "listed in BootOrder), re-run after the BootOrder reorder above has taken effect." >&2
      echo "Not rebooting -- purge did not fully succeed." >&2
      exit 1
    fi
    echo "All non-PXE entries deleted successfully."
    do_boot_pxe
    ;;
  boot-pxe)
    confirm "Reboot now to set next boot to PXE? This will wipe and reinstall the host."
    do_boot_pxe
    ;;
esac
