#!/bin/bash
ID=rosa
GRUBDIR='/boot/grub2'
GRUBENV="${GRUBDIR}/grubenv"
MNTPOINT='/boot/efi'
DEBUG='no'
FORCE='no'

if [ "$1" ]; then
    DEBUG=yes
    [ "$1" = '--force' ] && FORCE='yes'
fi

echo_exit() {
    [ "$DEBUG" != 'no' -o "$2" -ne 0 ] && echo "$1"
    exit "$2"
}

# Continue only if bootloaders were installed by Rosa
[ "$(find /boot -name 'rosa-iso-build*')" ] || echo_exit "File rosa-iso-build not found. Aborting $0" 0

# Continue if GRUB directory exists
[ -d "$GRUBDIR" ] || echo_exit "$GRUBDIR not found. Aborting $0" ${LINENO}

# If live OS, no need to install bootloader.
# Anaconda installs GRUB earlier.
if losetup | grep -Eq 'LiveOS/rootfs.img|layer-base/0'; then
    echo "Rosa live detected, skip grub2-install"
else
    # 1. Reinstall bootloader (works for both EFI and BIOS)
    # Determine the mount point for the boot device
    BOOT_MOUNT="/boot"
    [ -d /sys/firmware/efi ] && BOOT_MOUNT="/boot/efi"

    if [ -d "$BOOT_MOUNT" ]; then
        # Remove Btrfs subvolume suffix like [/root] or [/@]
        BOOT_PART=$(findmnt -n -o SOURCE -T "$BOOT_MOUNT" 2>/dev/null | sed 's/\[.*\]//')
        if [ -n "$BOOT_PART" ]; then
            BOOT_DISK_NAME=$(lsblk -no pkname "$BOOT_PART" 2>/dev/null)
            if [ -n "$BOOT_DISK_NAME" ]; then
                BOOT_DEV="/dev/$BOOT_DISK_NAME"
            else
                BOOT_DEV="$BOOT_PART"
            fi
    
            [ "$DEBUG" != 'no' ] && echo "Installing grub to $BOOT_DEV (from $BOOT_MOUNT)"
            grub2-install "$BOOT_DEV" || echo_exit "grub2-install failed for $BOOT_DEV" ${LINENO}
        else
            echo_exit "$BOOT_MOUNT - not found" ${LINENO}
        fi
    fi
fi

# 2. Setup grubenv (common part, works on any filesystem)
if [ ! -f "${GRUBENV}" ]; then
    grub2-editenv "${GRUBENV}" create || echo_exit "Create ${GRUBENV} - failed!" ${LINENO}
fi

FSTYPE=$(stat -f -c %T "${GRUBDIR}" 2>/dev/null || echo "unknown")
if [ "${FSTYPE}" = "btrfs" ]; then
    EDITENV_VER=$(grub2-editenv -V 2>/dev/null | awk '{print $NF}')
    
    # Only proceed if grub2-editenv version is >= 2.16
    if [ -n "${EDITENV_VER}" ] && [ "$(printf '%s\n' "2.16" "${EDITENV_VER}" | sort -V | head -n1)" = "2.16" ]; then
        if ! grub2-editenv "${GRUBENV}" list | grep -q "^env_block="; then
            if grub2-editenv "${GRUBENV}" set "env_block=512+1" 2>/dev/null ; then
                # need to enable savedefault on btrfs
                update-grub2  || echo_exit "update-grub2 failed after setting env_block" ${LINENO}
            else
                echo "Set env_block in ${GRUBENV} - failed!"
            fi
        fi
    else
        [ "$DEBUG" != 'no' ] && echo "grub2-editenv version ${EDITENV_VER} is too old (need >= 2.16). Skipping env_block setup."
    fi
fi

# 3. EFI-specific part (only if the system is booted in UEFI mode)
if [ ! -d /sys/firmware/efi ]; then
    [ "$DEBUG" != 'no' ] && echo "UEFI not detected. Skipping EFI bootloader installation."
    exit 0
fi

# Set EFI bootloader file prefix (PREFIX=x64 -> file bootx64.efi)
# According to UEFI spec, x64 is in small letters, others are in big letters
ARCH=$(arch)
case "$ARCH" in
    'x86_64')  PREFIX=x64 ;;
    'aarch64' | 'arm64') PREFIX=AA64 ;;
    'armhf')   PREFIX=ARM ;;
    'i386'|'i586'|'i686') PREFIX=IA32 ;;
    'riscv64') PREFIX=RISCV64 ;;
    'loongarch64') PREFIX=LOONGARCH64 ;;
    *)  echo_exit "Unknown architecture: $ARCH" ${LINENO} ;;
esac

# Check if EFI partition is mounted
if [ ! -d "$MNTPOINT" ]; then
    echo_exit "$MNTPOINT not found. Aborting $0" ${LINENO}
fi

pushd $MNTPOINT >/dev/null
EFIDIR="$(ls -1d {EFI,efi,Efi} 2>/dev/null | tail -n1)"
[ "$EFIDIR" ] || EFIDIR="EFI"
mkdir -p ${EFIDIR}
pushd ${EFIDIR} >/dev/null
df -t vfat ./ 2>/dev/null 1>&2 || exit ${LINENO}
[ -d "$ID" ] || echo_exit "${EFIDIR}/$ID not found. Aborting $0" ${LINENO}

BOOTDIR=$(ls -1d {BOOT,boot,Boot} 2>/dev/null | tail -n1)
[ "$BOOTDIR" ] || BOOTDIR='BOOT'
mkdir -p $BOOTDIR
if [ "$FORCE" != 'yes' ]; then
    # Continue if /EFI/BOOT is empty, or if Rosa bootloaders are present
    [ "$(ls -1 $BOOTDIR)" ] && [ ! -f "$BOOTDIR/${ID}.readme" ] && \
    echo_exit "${EFIDIR}/${BOOTDIR} must be empty or contain Rosa bootloaders. Aborting $0" ${LINENO}
fi

rmdir ${BOOTDIR} 2>/dev/null || (rm -rf ${BOOTDIR}.bak ; mv -f ${BOOTDIR} ${BOOTDIR}.bak)
mkdir -p $BOOTDIR
cp -f ${ID}/{grub${PREFIX}.efi,mm${PREFIX}.efi} ${BOOTDIR}/  2>/dev/null 
if [ -f ${ID}/shim${PREFIX}.efi ]; then 
    cp -f ${ID}/shim${PREFIX}.efi ${BOOTDIR}/BOOT${PREFIX}.efi
elif [ -f ${ID}/grub${PREFIX}.efi ]; then
    cp -f ${ID}/grub${PREFIX}.efi ${BOOTDIR}/BOOT${PREFIX}.efi
else
    echo_exit "Efi bootloader not found, need shim${PREFIX}.efi or grub${PREFIX}.efi" ${LINENO}
fi
cat << EOF > "${BOOTDIR}/${ID}.readme"
Rosa efi boot hook. Do not remove this file.
Files in ${EFIDIR}/${BOOTDIR}/  were copied from the ${EFIDIR}/${ID}/
EOF
popd >/dev/null
popd >/dev/null

exit 0
