#!/bin/bash

# --- SAFETY CHECKS ---
if [[ $EUID -ne 0 ]]; then
   echo "ERROR: This script must be run as root (or inside arch-chroot)."
   exit 1
fi

if [ ! -f "/etc/arch-release" ]; then
    echo "ERROR: Run this INSIDE the arch-chroot environment!"
    exit 1
fi

echo "================================================="
echo "    ARCHINSTALL UKI + SHIM INSTALLER"
echo "================================================="

# --- 1. PARTITION SELECTION ---
lsblk -p
echo ""

# Auto-detect root
ROOT_DETECT=$(findmnt -n -o SOURCE /)
echo "[*] Detected Root: $ROOT_DETECT"

# Use < /dev/tty to allow input during a curl | bash session
read -p "Confirm ROOT partition (default: $ROOT_DETECT): " ROOT_PART < /dev/tty
ROOT_PART=${ROOT_PART:-$ROOT_DETECT}

if [ ! -b "$ROOT_PART" ]; then
    echo "ERROR: $ROOT_PART is not a valid block device. Exiting."
    exit 1
fi

read -p "Enter Ventoy USB partition (e.g., /dev/sda1): " V_PART < /dev/tty
if [ ! -b "$V_PART" ]; then
    echo "ERROR: $V_PART is not a valid block device. Exiting."
    exit 1
fi

# --- 2. PURGE LEFTOVERS ---
echo "[*] Purging conflicting bootloaders..."
rm -rf /boot/EFI/arch /boot/EFI/systemd /boot/EFI/grub /boot/loader
rm -f /boot/EFI/BOOT/BOOTX64.EFI 2>/dev/null

echo "[*] Clearing NVRAM (efibootmgr)..."
# Using -b and -B to clear existing entries
for entry in $(efibootmgr | grep -iE 'arch|grub|systemd|linux boot' | awk '{print $1}' | sed 's/Boot//' | sed 's/\*//'); do
    echo "[*] Removing entry $entry"
    efibootmgr -b "$entry" -B > /dev/null
done

# --- 3. UKI ARCHITECTURE ---
echo "[*] Configuring Unified Kernel Image parameters..."
mkdir -p /etc/kernel
echo "root=$ROOT_PART rw" > /etc/kernel/cmdline

PRESET_FILE=$(ls /etc/mkinitcpio.d/*.preset | head -n 1)
if [ -z "$PRESET_FILE" ]; then
    echo "ERROR: No mkinitcpio preset found. Did you install a kernel?"
    exit 1
fi

echo "[*] Modifying $PRESET_FILE for UKI..."
# Ensure the EFI/BOOT directory exists before mkinitcpio runs
mkdir -p /boot/EFI/BOOT
sed -i 's|^#*default_uki=.*|default_uki="/boot/EFI/BOOT/grub.efi"|' "$PRESET_FILE"
sed -i 's|^default_image=.*|#default_image="/boot/initramfs-linux.img"|' "$PRESET_FILE"

# --- 4. THE HEIST ---
echo "[*] Generating UKI (grub.efi)..."
mkinitcpio -P

echo "[*] Mounting Ventoy to retrieve Secure Boot Shim..."
mkdir -p /mnt_ventoy
if mount "$V_PART" /mnt_ventoy; then
    cp /mnt_ventoy/EFI/BOOT/BOOTX64.EFI /boot/EFI/BOOT/BOOTX64.EFI
    cp /mnt_ventoy/EFI/BOOT/MokManager.efi /boot/EFI/BOOT/MokManager.efi
    umount /mnt_ventoy
    echo "[*] Shim successfully moved to /boot/EFI/BOOT/"
else
    echo "ERROR: Could not mount Ventoy partition. Check the device name."
    exit 1
fi
rmdir /mnt_ventoy

# --- 5. FINAL REPORT ---
echo "================================================="
echo "   INSTALL SUCCESSFUL"
echo "================================================="
echo " FILES IN BOOT DIRECTORY:"
ls -lh /boot/EFI/BOOT/
echo "-------------------------------------------------"
echo " 1. Type: exit"
echo " 2. Type: umount -R /mnt"
echo " 3. REBOOT and PULL THE USB DRIVE."
echo " 4. When the MokManager screen appears:"
echo "    Enroll Hash -> Internal Drive -> EFI -> BOOT -> grub.efi"
echo "================================================="