#!/usr/bin/env bash
# Create a GRUB USB disk that boots on both BIOS (legacy) and UEFI.
# It copies kernel (vmlinuz) and initramfs, installs GRUB for i386-pc and x86_64-efi,
# and writes a minimal grub.cfg with optional kernel parameters.
#
# DANGEROUS: This script will repartition and format the target device (e.g., /dev/sdX).
#
# Usage:
#   sudo scripts/make-grub-usb.sh /dev/sdX \
#        --kernel dist/vmlinuz.efi \
#        --initrd dist/initramfs.cpio.xz \
#        --kparams "console=ttyS0 initdebug=true" \
#        --label ZOSBOOT \
#        --no-confirm
#
# Defaults (resolved relative to repo root):
#   --kernel  dist/vmlinuz.efi
#   --initrd  dist/initramfs.cpio.xz
#   --label   ZOS
#   --kparams ""  (none)
#
# Requirements on host:
#   - grub-install (supports --target=i386-pc and --target=x86_64-efi)
#   - parted, sfdisk (optional), dosfstools (mkfs.vfat), util-linux (lsblk, partprobe)
#   - run as root
#
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

DEVICE="${1:-}"
shift || true

KERNEL="${PROJECT_ROOT}/dist/vmlinuz.efi"
# Default: no separate initrd because initramfs is embedded in kernel. Enable with --with-initrd or --initrd.
INITRD=""
KPARAMS=""
LABEL="ZOS"
NO_CONFIRM=0
ESP_SIZE_MB=512
# Request including separate initrd
WANT_INITRD=0

error() { echo "[ERROR] $*" >&2; }
info()  { echo "[INFO]  $*"; }
warn()  { echo "[WARN]  $*"; }
die()   { error "$*"; exit 1; }

require_root() {
  [[ "$EUID" -eq 0 ]] || die "This script must be run as root."
}

command_exists() { command -v "$1" &>/dev/null; }

parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --kernel)
        KERNEL="$(realpath -m "${2:-}")"; shift 2 ;;
      --initrd)
        INITRD="$(realpath -m "${2:-}")"; shift 2 ;;
      --with-initrd)
        WANT_INITRD=1; shift ;;
      --kparams)
        KPARAMS="${2:-}"; shift 2 ;;
      --label)
        LABEL="${2:-ZOS}"; shift 2 ;;
      --no-confirm)
        NO_CONFIRM=1; shift ;;
      --esp-size-mb)
        ESP_SIZE_MB="${2:-512}"; shift 2 ;;
      -*)
        die "Unknown option: $1" ;;
      *)
        die "Unexpected argument: $1" ;;
    esac
  done
}

confirm_dangerous() {
  info "About to wipe and repartition device: ${DEVICE}"
  lsblk -o NAME,SIZE,TYPE,MOUNTPOINT "${DEVICE}"
  if [[ $NO_CONFIRM -eq 1 ]]; then
    info "--no-confirm provided; proceeding without interactive confirmation"
    return
  fi
  echo
  read -r -p "Type the DEVICE path to CONFIRM (e.g., ${DEVICE}) or 'abort' to cancel: " ans
  [[ "$ans" == "$DEVICE" ]] || die "Confirmation mismatch or aborted."
}

check_prereqs() {
  command_exists parted      || die "parted not found"
  command_exists mkfs.vfat   || die "mkfs.vfat (dosfstools) not found"
  command_exists grub-install || die "grub-install not found"
  command_exists lsblk       || die "lsblk not found"
  command_exists partprobe   || die "partprobe not found"
}

resolve_defaults() {
  [[ -b "$DEVICE" ]] || die "Device not found or not a block device: ${DEVICE}"
  [[ -f "$KERNEL" ]] || die "Kernel file not found: ${KERNEL}"
  if [[ $WANT_INITRD -eq 1 ]]; then
    # If user asked for separate initrd but none specified, use default path
    if [[ -z "${INITRD}" ]]; then
      INITRD="${PROJECT_ROOT}/dist/initramfs.cpio.xz"
    fi
    [[ -f "$INITRD" ]] || die "Requested --with-initrd but initramfs not found: ${INITRD}"
  else
    # By default, no separate initrd (initramfs is embedded in kernel)
    INITRD=""
  fi
}

umount_partitions() {
  info "Unmounting any mounted partitions of ${DEVICE}"
  local kids
  kids="$(lsblk -nr -o NAME "${DEVICE}" | tail -n +2 || true)"
  if [[ -n "${kids}" ]]; then
    while read -r name; do
      [[ -z "$name" ]] && continue
      local path="/dev/${name}"
      if mountpoint -q -- "/dev/${name}" 2>/dev/null; then
        umount -f "/dev/${name}" || true
      fi
      # also try mounts by path
      local mp
      mp="$(lsblk -nr -o MOUNTPOINT "$path" || true)"
      if [[ -n "$mp" ]]; then
        for m in $mp; do
          [[ -n "$m" ]] && umount -f "$m" || true
        done
      fi
    done <<< "${kids}"
  fi
}

partition_device_gpt() {
  info "Creating GPT with BIOS boot + ESP on ${DEVICE}"
  # Create a new GPT
  parted -s "${DEVICE}" mklabel gpt

  # Create BIOS boot partition 1MiB..3MiB
  parted -s "${DEVICE}" mkpart biosgrub 1MiB 3MiB
  parted -s "${DEVICE}" set 1 bios_grub on

  # Create ESP (FAT32) from 3MiB to 3MiB + ESP_SIZE_MB
  local esp_end="$((3 + ESP_SIZE_MB))MiB"
  parted -s "${DEVICE}" mkpart esp fat32 3MiB "${esp_end}"
  parted -s "${DEVICE}" set 2 esp on

  # Inform kernel
  partprobe "${DEVICE}"
  sleep 1
}

format_esp() {
  local part="${DEVICE}2"
  [[ -b "$part" ]] || die "ESP partition not found: ${part}"
  info "Formatting ESP ${part} as FAT32 (label=${LABEL})"
  mkfs.vfat -F32 -n "${LABEL}" "$part"
}

mount_esp() {
  ESP_MNT="$(mktemp -d)"
  info "Mounting ESP at ${ESP_MNT}"
  mount "${DEVICE}2" "${ESP_MNT}"
}

install_grub() {
  info "Installing GRUB (BIOS i386-pc) to ${DEVICE}"
  grub-install --target=i386-pc --boot-directory="${ESP_MNT}/boot" --recheck "${DEVICE}"

  info "Installing GRUB (UEFI x86_64-efi) to ESP"
  mkdir -p "${ESP_MNT}/EFI/BOOT"
  grub-install --target=x86_64-efi --efi-directory="${ESP_MNT}" --boot-directory="${ESP_MNT}/boot" --removable --recheck
}

copy_kernel_initrd() {
  mkdir -p "${ESP_MNT}/boot"
  info "Copying kernel to ESP /boot/vmlinuz"
  cp -f "${KERNEL}" "${ESP_MNT}/boot/vmlinuz"
  if [[ -n "${INITRD}" ]]; then
    info "Copying initramfs to ESP /boot/initramfs.cpio.xz"
    cp -f "${INITRD}" "${ESP_MNT}/boot/initramfs.cpio.xz"
  fi
}

write_grub_cfg() {
  local cfg="${ESP_MNT}/boot/grub/grub.cfg"
  info "Writing GRUB config: ${cfg}"
  mkdir -p "$(dirname "$cfg")"
  cat > "$cfg" <<'EOF'
set default=0
set timeout=3

if [ "${grub_platform}" = "efi" ]; then
  set gfxpayload=keep
fi

menuentry "Zero-OS" {
  linux /boot/vmlinuz __KPARAMS__
__INITRD_LINE__
}
EOF

  local initrd_line=""
  if [[ -n "${INITRD}" ]]; then
    initrd_line="  initrd /boot/initramfs.cpio.xz"
  fi
  # Substitute placeholders
  sed -i \
    -e "s|__KPARAMS__|${KPARAMS}|g" \
    -e "s|__INITRD_LINE__|${initrd_line}|g" \
    "$cfg"

  # For removable media on UEFI, ensure fallback bootloader path exists
  if [[ -f "/usr/lib/grub/x86_64-efi/monolithic/grubx64.efi" ]]; then
    cp -f "/usr/lib/grub/x86_64-efi/monolithic/grubx64.efi" "${ESP_MNT}/EFI/BOOT/BOOTX64.EFI" || true
  elif [[ -f "${ESP_MNT}/EFI/BOOT/grubx64.efi" ]]; then
    cp -f "${ESP_MNT}/EFI/BOOT/grubx64.efi" "${ESP_MNT}/EFI/BOOT/BOOTX64.EFI" || true
  fi
}

cleanup() {
  set +e
  if mountpoint -q "${ESP_MNT}" 2>/dev/null; then
    info "Unmounting ESP ${ESP_MNT}"
    umount "${ESP_MNT}"
  fi
  [[ -n "${ESP_MNT:-}" ]] && rmdir "${ESP_MNT}" 2>/dev/null || true
}

main() {
  require_root
  [[ -n "${DEVICE}" ]] || die "Usage: $0 /dev/sdX [--kernel path] [--initrd path] [--kparams \"...\"] [--label ZOS] [--no-confirm]"
  parse_args "$@"
  resolve_defaults
  check_prereqs

  confirm_dangerous
  trap cleanup EXIT INT TERM

  umount_partitions
  partition_device_gpt
  format_esp
  mount_esp
  install_grub
  copy_kernel_initrd
  write_grub_cfg

  info "Done. You can now boot from the USB on BIOS and UEFI systems."
}

main "$@"