#!/bin/bash
# enter the zoid (zfs on root with zraid) - developed by acidvegas (https://git.acid.vegas/void)
#      boot: https://github.com/leahneukirchen/hrmpf
# reference: https://docs.zfsbootmenu.org/en/v2.3.x/guides/void-linux/uefi.html (do we need to make any updates?)

set -xev

# Configuration
export HOSTNAME=blackhole
export BOOT_DRIVE=/dev/sde # Use the internal USB drive for the boot partition
export BOOT_METHOD=direct # Use direct or refind
export POOL_DRIVES="/dev/sda /dev/sdb /dev/sdc /dev/sdd" # Verify these with lsblk before running
export RAIDZ_PARITY="1" # Number of drives to use for the RAID-Z parity (must be 1 or greater otherwise why are you using ZFS?)


checks() {
	# Check if the system is using UEFI or BIOS
	if [ ! -d /sys/firmware/efi ]; then
		echo "System must be using UEFI"
		exit 1
	fi

	# Check if all drives exist and are valid
	for d in $BOOT_DRIVE $POOL_DRIVES; do
		if [ ! -b $d ]; then
			echo "Drive $d does not exist"
			exit 1
		fi
	done

	# Check if the boot method is valid
	if [ $BOOT_METHOD != "direct" ] && [ $BOOT_METHOD != "refind" ]; then
        echo "Boot method must be direct or refind"
        exit 1
    fi

	# Check if the RAID-Z parity is valid
	if [ $RAIDZ_PARITY -lt 1 ]; then
		echo "RAID-Z parity must be 1 or greater"
		exit 1
	fi
}


setup_zfs() {
    # Validation
    check

    # Generate the hostid
	source /etc/os-release
	export ID
	zgenhostid -f 0x00bab10c

	# Prepare the boot drive
	wipefs -a $BOOT_DRIVE
	sgdisk --zap-all $BOOT_DRIVE
	sgdisk -n "1:1m:+1g" -t "1:ef00" $BOOT_DRIVE

	# Prepare the ZFS pool drives
	for d in $POOL_DRIVES; do
		wipefs -a $d
		sgdisk --zap-all $d
		sgdisk -n "1:0:-10m" -t "1:bf00" "$d"
		if zdb -l "$d" &> /dev/null; then
			zpool labelclear -f "$d"
		fi
	done

	# Create the ZFS pool (should we use noatime=on instead of relatime=on?)
	ZFS_POOL_DRIVES=$(echo $(for dev in $POOL_DRIVES; do find /dev/disk/by-id/ -samefile $(readlink -f "$dev") ! -name "*-part*" -print -quit; done))
	zpool create -f -o ashift=12 -O compression=lz4 -O acltype=posixacl -O xattr=sa -O relatime=on -o autotrim=on -o compatibility=openzfs-2.1-linux -m none zroot raidz${RAIDZ_PARITY} $ZFS_POOL_DRIVES

	# Create the ZFS datasets
	zfs create -o mountpoint=none zroot/ROOT
	zfs create -o mountpoint=/ -o canmount=noauto zroot/ROOT/$ID
	zfs create -o mountpoint=/home zroot/home
	zpool set bootfs=zroot/ROOT/$ID zroot

	# Export and import the ZFS pool
	zpool export zroot
	zpool import -N -R /mnt zroot
	zfs mount zroot/ROOT/$ID
	zfs mount zroot/home

	# Trigger udev
	udevadm trigger

	# Install base system
	XBPS_ARCH=x86_64 xbps-install -S -R https://mirrors.servercentral.com/voidlinux/current -r /mnt base-system

	# Copy the hostid into the new system
	cp /etc/hostid /mnt/etc

	# Chroot into the new system
	echo "entering the void..."
	xchroot /mnt
}


setup_chroot() {
    # Set the root password
	echo "root:root" | chpasswd

	# Update the package manager and install some basic packages
	xbps-install -Suy nano wget

	# Install the non-free repository
	xbps-install -y void-repo-nonfree
	xbps-install -Suy

	# Install & enable the intel microcode service
	xbps-install -y intel-ucode
	ln -sf /etc/sv/intel-ucode /etc/runit/runsvdir/default/

	# Set the timezone & hardware clock
	ln -sf /usr/share/zoneinfo/America/New_York /etc/localtime
	hwclock --systohc

	# Set the hostname
	echo "$HOSTNAME" > /etc/hostname

	# Set the rc.conf variables
	printf "HOSTNAME=\"$HOSTNAME\"\nHARDWARECLOCK=\"UTC\"\nTIMEZONE=\"America/New_York\"\nKEYMAP=us\n" > /etc/rc.conf

	# Set nameservers
	printf "nameserver 9.9.9.9" > /etc/resolv.conf

	# Set the locales
	printf "en_US.UTF-8 UTF-8\nen_US ISO-8859-1\n" > /etc/default/libc-locales
	xbps-reconfigure -f glibc-locales

	# Set the dracut configuration
	printf "nofsck=\"yes\"\nadd_dracutmodules+=\" zfs \"\nomit_dracutmodules+=\" btrfs \"\n" > /etc/dracut.conf.d/zol.conf

	# Install the zfs package
	xbps-install -y zfs

	# Set the zfsbootmenu command line options
	zfs set org.zfsbootmenu:commandline="quiet loglevel=4" zroot/ROOT

	# Setup & mount the boot partition
	mkfs.vfat -F32 ${BOOT_DRIVE}1
	BOOT_UUID=$(blkid -s UUID -o value ${BOOT_DRIVE}1)
	echo "UUID=$BOOT_UUID /boot/efi vfat defaults 0 0" > /etc/fstab
	mkdir -p /boot/efi
	mount /boot/efi

	# Install and setup zfsbootmenu
	xbps-install -S zfsbootmenu gummiboot-efistub yq
	yq -iy '.Global.ManageImages=true | .Global.BootMountPoint="/boot/efi" | .Components.Enabled=false | .EFI.ImageDir="/boot/efi/EFI/zbm" | .EFI.Versions=false | .EFI.Enabled=true | .Kernel.CommandLine="quiet loglevel=0"' /etc/zfsbootmenu/config.yaml
	generate-zbm

	# Apply boot method
	# Note      : Some systems can have issues with EFI boot entries, you might need to use a well-known EFI file name.
	# Reference : https://docs.zfsbootmenu.org/en/v2.3.x/general/portable.html
	if [ $BOOT_METHOD == "direct" ]; then
	    xbps-install efibootmgr
		efibootmgr -c -d "$BOOT_DRIVE" -p "1" -L "ZFSBootMenu (Backup)" -l '\EFI\ZBM\VMLINUZ-BACKUP.EFI'
		efibootmgr -c -d "$BOOT_DRIVE" -p "1" -L "ZFSBootMenu" -l '\EFI\ZBM\VMLINUZ.EFI'
	elif [ $BOOT_METHOD == "refind" ]; then
    	xbps-install -y refind
        refind-install
    	rm /boot/refind_linux.conf
    	printf "\"Boot default\"  \"quiet loglevel=0 zbm.skip\"\n\"Boot to menu\"  \"quiet loglevel=0 zbm.show\"\n" > /boot/efi/EFI/ZBM/refind_linux.conf
    	# Note     : Everything below this line is a "hacky" solution to a problem I was having with the zfsbootmenu package
    	# Reference: https://github.com/zbm-dev/zfsbootmenu/issues/293
    	mkdir -p /boot/efi/EFI/BOOT
    	mvrefind /boot/efi/EFI/refind /boot/efi/EFI/BOOT
    	temp=$(mktemp -d)
    	wget -O $temp/latest.tar.gz https://get.zfsbootmenu.org/latest.tar.gz
    	tar xvf $temp/latest.tar.gz -C $temp/
    	rm $temp/latest.tar.gz
    	mv $temp/zfs*/* /boot/efi/EFI/ZBM/
    	rm /boot/efi/EFI/ZBM/vmlinuz.efi
    	xbps-remove zfsbootmenu
    fi

	# Reconfigure the system
	xbps-reconfigure -fa

	# Exit the chroot environment
	echo "exiting the void..."
	exit
}



# Check the command
if [ "$#" -ne 1 ]; then
	echo "usage: $0 [zfs|chroot|final]"
	exit 1
fi

# Execute the command
case "$1" in
	zfs)    setup_zfs ;;
	chroot) setup_chroot ;;
	final)  umount -n -R /mnt; zpool export zroot; reboot ;;
	*)      echo "usage: $0 [zfs|chroot|final]" && exit 1 ;;
esac