#!/bin/bash

set -euo pipefail

# ==============================================================
#  ZFSBootMenu Deployment Script (Legacy Boot / Devuan)
#  Safe Unified Kernel Image Extractor & GRUB Auto Integration
# ==============================================================

# Boot Flow Visualization
# 1.    BIOS/Legacy: Execution starts at the MBR/VBR.
# 2.    GRUB Stage 1 & 2: GRUB loads its modules and reads your custom script from /etc/grub.d/.
# 3.    The "Direct" Load: GRUB uses the linux and initrd commands to pull the extracted ZBM components into RAM.
# 4.    ZBM Environment: ZBM initializes, finds your ZFS datasets, and provides the UI.
# 5.    kexec: Once you select a kernel in ZBM, it uses kexec to replace itself with your actual Devuan kernel.

# --- CLI Argument Parsing ---
FORCE=0
for arg in "$@"; do
  case "$arg" in
  --force | --force=1 | --force=true)
    FORCE=1
    ;;
  esac
done

#
# --- Configuration ---
REPO="zbm-dev/zfsbootmenu"
ZBM_DIR="/boot/zfsbootmenu"
KEYRING="/usr/share/keyrings/zfsbootmenu.gpg"
PUBKEY_URL="https://raw.githubusercontent.com/${REPO}/master/releng/keys/zfsbootmenu.pub.gpg"
HASH_FILE="sha256.txt"
SIG_FILE="sha256.sig"

# --- 1. Dependency Check ---
echo "[*] Checking dependencies..."
for cmd in objcopy curl jq gpg findmnt zfs; do
  if ! command -v "$cmd" >/dev/null 2>&1; then
    echo "Installing missing dependency: $cmd..."
    apt-get update -qq && apt-get install -y binutils curl jq gnupg zfsutils-linux >/dev/null
  fi
done

# --- 2. Failure Recovery Logic ---
restore_backups() {
  echo "!!! Error encountered. Restoring previous version from .old backups..."
  for file in vmlinuz-zbm initramfs-zbm.img; do
    if [ -f "$ZBM_DIR/$file.old" ]; then
      mv -f "$ZBM_DIR/$file.old" "$ZBM_DIR/$file"
    fi
  done
}

# --- 3. Version & Hash Idempotency ---
echo "[*] Checking for latest ZFSBootMenu release..."
LATEST_TAG=$(curl -sfSL "https://api.github.com/repos/$REPO/releases/latest" | jq -r .tag_name)
# LATEST_TAG=$(curl -sfSL "https://api.github.com/repos/$REPO/releases/latest" | jq -r '.tag_name' | tr -d '"')
echo "Using clean tag: '$LATEST_TAG'" # v3.1.0 ✓

MACHINE="$(uname -m)"                            # x86_64
KERNEL="$(uname -r)"                             # 6.12.6-amd64
KERNEL_VERSION=$(echo "$KERNEL" | cut -d. -f1,2) # 6.12
KERNEL_NAME="$(uname -s)"                        # Linux
KERNEL_NAME="${KERNEL_NAME,,}"                   # linux lowercase
EFI_NAME="zfsbootmenu-recovery-${MACHINE}-${LATEST_TAG}-${KERNEL_NAME}${KERNEL_VERSION}.EFI"
HASH_URL="https://github.com/$REPO/releases/download/$LATEST_TAG/${HASH_FILE}"
SIG_URL="https://github.com/$REPO/releases/download/$LATEST_TAG/${SIG_FILE}"
REMOTE_HASH_DATA=$(curl -sfSL --connect-timeout 10 "$HASH_URL")
# REMOTE_HASH=$(echo "$REMOTE_HASH_DATA" | grep "$EFI_NAME" | awk '{print $1}')
# REMOTE_HASH=$(echo "$REMOTE_HASH_DATA" | grep "$EFI_NAME" | sed -n "s/.*= \([0-9a-f]*\).*/\1/p")
REMOTE_HASH=$(echo "$REMOTE_HASH_DATA" | grep "$EFI_NAME" | sed -n "s/.*= *\([0-9a-f]\{64\}\).*/\1/p")

if [[ -f "$ZBM_DIR/.hash" && "$(cat "$ZBM_DIR/.hash")" == "$REMOTE_HASH" && "$FORCE" -ne 1 ]]; then
  echo "✔ ZFSBootMenu $LATEST_TAG already installed and verified. Exiting."
  exit 0
fi

# --- 4. Workspace & Error Handling ---
TMP_DIR=$(mktemp -d)

cleanup_on_error() {
  local exit_code=$?
  if [[ $exit_code -ne 0 ]]; then
    echo "ERROR: Script failed (line $BASH_LINENO). Rolling back..."
    restore_backups
  fi
  rm -rf "$TMP_DIR"
  exit "$exit_code"
}
trap 'cleanup_on_error' EXIT ERR

# --- 5. Secure Download & Verification ---
echo "[*] Ensuring GPG keyring exists..."
[ ! -f "$KEYRING" ] && sudo curl -sfSL "$PUBKEY_URL" -o "$KEYRING"

cd "$TMP_DIR"
echo "[*] Downloading latest UKI: $EFI_NAME..."
curl -sfSLO "https://github.com/$REPO/releases/download/$LATEST_TAG/$EFI_NAME"
curl -sfSLO "${HASH_URL}"

# echo "[*] Verifying GPG signature..."
# gpg --no-default-keyring --keyring "$KEYRING" --verify "${SIG_FILE}" "$EFI_NAME" 2>/dev/null || {
#   echo "ERROR: GPG signature verification failed!"
#   exit 1
# }

# CORRECT SHA256:
curl -sfSLO "$HASH_URL"
echo "$REMOTE_HASH  $EFI_NAME" | sha256sum -c || {
  echo "ERROR: SHA256 checksum failed!"
  exit 1
}

# --- 6. Extraction & Validation ---
echo "[*] Extracting components from UKI..."
objcopy --dump-section .linux="vmlinuz-zbm.tmp" "$EFI_NAME"
objcopy --dump-section .initrd="initramfs-zbm.img.tmp" "$EFI_NAME"

# Validate outputs
if [[ ! -s "vmlinuz-zbm.tmp" || ! -s "initramfs-zbm.img.tmp" ]]; then
  echo "ERROR: Extracted components are empty or missing! Abort."
  exit 1
fi

echo "[*] Replacing ZFSBootMenu components..."
sudo bash -c "
  [ -f "$ZBM_DIR/vmlinuz-zbm" ] && cp "$ZBM_DIR/vmlinuz-zbm" "$ZBM_DIR/vmlinuz-zbm.old"
  [ -f "$ZBM_DIR/initramfs-zbm.img" ] && cp "$ZBM_DIR/initramfs-zbm.img" "$ZBM_DIR/initramfs-zbm.img.old"

  mv vmlinuz-zbm.tmp "$ZBM_DIR/vmlinuz-zbm"
  mv initramfs-zbm.img.tmp "$ZBM_DIR/initramfs-zbm.img"
  chmod 600 "$ZBM_DIR/vmlinuz-zbm" "$ZBM_DIR/initramfs-zbm.img"
  echo $REMOTE_HASH >$ZBM_DIR/.hash
"

# --- 7. GRUB Integration ---
ZBM_POOL=$(findmnt -n -o SOURCE -T "$ZBM_DIR" | cut -d'/' -f1)
ZBM_SOURCE=$(findmnt -n -o SOURCE -T "$ZBM_DIR") # rpool/ROOT/devuan-1
ZBM_DATASET="${ZBM_SOURCE#*/}"                   # ROOT/devuan-1
ZBM_DATASET_PATH="/${ZBM_DATASET}@/"             # /ROOT/devuan-1@/

# Combine dataset location with the ZBM storage path
REL_PATH="${ZBM_DATASET_PATH}${ZBM_DIR#/}" # /ROOT/devuan-1@/boot/zfsbootmenu

echo "[*] Generating GRUB configuration for pool: $ZBM_POOL..."

# | Parameter                   | Purpose                                            | Notes                                                             |
# | --------------------------- | -------------------------------------------------- | ----------------------------------------------------------------- |
# | zbm.prefer=<pool>           | Tells ZBM which ZFS pool to import first           | Highly recommended if you have multiple zpools or bootable clones |
# | zbm.import_delay=<seconds>  | Adds a startup delay before auto-importing pools   | Useful for slow disks or sparse device init                       |
# | zbm.readonly=1              | Forces all imports read‑only                       | For emergency/recovery boot entries                               |
# | zbm.timeout=<sec>           | Default menu timeout                               | Optional override of ZBM’s own internal timeout                   |
# | zbm.skip=<boolean>          | Skip pool discovery, assume you’ll import manually | Optional, advanced                                                |
# | loglevel=<n>                | Controls kernel verbosity                          | 4 is a balanced choice; 7 for deep debugging                      |
# | rd.vconsole.keymap=<keymap> | Keyboard layout for early TTY                      | Optional, good for non‑US systems                                 |
# | rd.vconsole.font=<font>     | Sets framebuffer console font (e.g. ter-124n)      | Helps when small fonts are hard to read on hi‑res screens         |

# Define arguments to add
# Font sizes     : ter-v32n ter-v28n ter-v24n ter-v20n ter-v14n
# Font sizes bold: ter-v32b ter-v28b ter-v24b ter-v20b ter-v14b

# FIXME: If this where placed in /etc/default/zfsbootmenu/ it could be sourced here rather than hardcoded here. Much like /etc/default/grub
# The scripts is separated from the settings. Maybe something I would do if I were packaging zbm.
GRUB_DEFAULTS='loglevel=4 zbm.import_delay=5 video=vesafb:1920x1200-32@60 rd.vconsole.keymap=uk'

echo "DEBUG: ZBM_SOURCE='$ZBM_SOURCE'"
echo "DEBUG: ZBM_DATASET='$ZBM_DATASET'"
echo "DEBUG: ZBM_DATASET_PATH='$ZBM_DATASET_PATH'"
echo "DEBUG: REL_PATH='$REL_PATH'"
echo "DEBUG: ZBM_DIR='$ZBM_DIR'"

conf_print_grub_menu_zbm() {
  cat <<EOF
#!/bin/sh
exec tail -n +3 \$0
# Dynamically generated by ZBM Deployment Script

submenu 'ZFSBootMenu>Recovery Options' {
  menuentry 'ZFSBootMenu (Direct Boot)' --class zfs --class gnu-linux {
      insmod part_gpt
      insmod zfs
      search --no-floppy --set=root --label $ZBM_POOL
      linux  $REL_PATH/vmlinuz-zbm ${GRUB_DEFAULTS} zbm.prefer=$ZBM_POOL
      initrd $REL_PATH/initramfs-zbm.img
  }

  menuentry 'ZFSBootMenu (Previous/Backup)' --class zfs --class gnu-linux {
      insmod part_gpt
      insmod zfs
      search --no-floppy --set=root --label $ZBM_POOL
      linux  $REL_PATH/vmlinuz-zbm.old  ${GRUB_DEFAULTS} zbm.prefer=$ZBM_POOL
      initrd $REL_PATH/initramfs-zbm.img.old
  }

  menuentry 'ZFSBootMenu (Recovery/Read-Only)' --class recovery --class zfs {
      insmod part_gpt
      insmod zfs
      search --no-floppy --set=root --label $ZBM_POOL
      linux  $REL_PATH/vmlinuz-zbm ${GRUB_DEFAULTS} zbm.readonly=1 zbm.prefer=$ZBM_POOL
      initrd $REL_PATH/initramfs-zbm.img
  }

  menuentry 'ZFSBootMenu (Recovery/Force Import)' --class recovery --class zfs {
      insmod part_gpt
      insmod zfs
      search --no-floppy --set=root --label $ZBM_POOL
      linux  $REL_PATH/vmlinuz-zbm ${GRUB_DEFAULTS} zbm.readonly=1 zbm.prefer=$ZBM_POOL!
      initrd $REL_PATH/initramfs-zbm.img
  }
}
EOF
}
conf_print_grub_menu_zbm | sudo tee /etc/grub.d/40_zfsbootmenu >/dev/null
sudo chmod +x /etc/grub.d/40_zfsbootmenu

if command -v update-grub >/dev/null 2>&1; then
  sudo update-grub
else
  sudo grub-mkconfig -o /boot/grub/grub.cfg
fi

echo "✔ Deployment complete, verify your grub.cfg entries below: "
grep -A5 'ZFSBootMenu' /boot/grub/grub.cfg