#!/bin/bash

if [[ "$(id -u)" -ne 0 ]]; then
	echo "ERROR: This needs to be run as root"
	exit 250
fi

declare -rA SUBVOLS_DEFAULT=(
  ["@home"]="home"
  ["@root"]="root"
  ["@srv"]="srv"
  ["@opt"]="opt"
  ["@local"]="usr/local"
  ["@cache"]="var/cache"
  ["@repos"]="/var/db/repos"
  ["@containers"]="var/lib/containers"
  ["@libvirt"]="var/lib/libvirt/images"
  ["@machines"]="var/lib/machines"
  ["@portables"]="var/lib/portables"
  ["@log"]="var/log"
  ["@spool"]="var/spool"
  ["@www"]="var/www"
  ["@snapshots"]=".snapshots"
)

function show_help() {
	echo "Usage: $0 <disk> [<options>...]"
	echo ""
	echo "Positional Arguments:"
	echo "<drive>   Disk device used for system"
	echo ""
	echo "Options:"
	echo "-h, --help        Help on this tool."
	echo "-e, --encryption  Enable LUKS encryption."
    echo "-c, --compression Enable BtrFS compression."
	echo "-s, --swap        Enable Swap/Hibernation support."
	echo "-d, --debug       Enable DEBUG mode for testing."
    echo ""
    echo "--stage <stage>   Installation using stagefile <stage>, for stage3 or stage4"
	echo "--mount           Mount partitions/subvolumes only. Useful for development or recovery"
	echo "--umount          Unmount partitions/subvolumes only"
    echo "--clean           Cleanup disk for clean slate"
	exit 0
}

function errormsg() {
	local err=$1
	shift

	echo "$*" >&2
	exit "$err"
}

function prepare_disk() {
	if [[ "$DEBUG" ]]; then
		local cmd="echo"
	else
		local cmd=""
	fi

    if [[ -b "$RootDisk" ]]; then
        if [[ -b "$EFIPart" || -b "$RootPart" ]]; then
            echo "FATAL: Disk already provisioned. Clean required to continue."
            exit 250
        fi

        if [[ "$ENCRYPTION" ]]; then
			if [[ "$UKI" ]]; then
				${cmd} parted --script --align=optimal -- "$RootDisk" \
					mklabel gpt \
					mkpart primary 1MiB 2048MiB \
					mkpart primary 2048MiB 3072MiB \
					mkpart primary 3072MiB -2048s \
					set 1 esp
			else
				${cmd} parted --script --align=optimal -- "$RootDisk" \
					mklabel gpt \
					mkpart primary 1MiB 100MiB \
					mkpart primary 100MiB 2048MiB \
					mkpart primary 2148MiB -2048s \
					set 1 esp
			fi
            #mkfs.vfat -F 32 -n "EFI" "$RootPart"
            #mkfs.ext4 -L "Boot" -m 0 "$EFIPart"
            #mkfs.btrfs -L "System" "$RootPart"
        else
			if [[ "$UKI" ]]; then
				${cmd} parted --script --align=optimal -- "$RootDisk" \
					mklabel gpt \
					mkpart primary 1MiB 2048MiB \
					mkpart primary 2048MiB -2048s \
					set 1 esp
			else
				${cmd} parted --script --align=optimal -- "$RootDisk" \
					mklabel gpt \
					mkpart primary 1MiB 100MiB \
					mkpart primary 100MiB -2048s \
					set 1 esp
			fi
            #mkfs.vfat -F 32 -n "EFI" "$RootPart"
            #mkfs.btrfs -L "System" "$RootPart"
        fi
	fi
}

function create_luks() {
	if [[ "$DEBUG" ]]; then
		local cmd="echo"
	else
		local cmd=""
	fi

	${cmd} cryptsetup --batch-mode --cipher aes-xts-plain64 --hash sha512 --use-random --verify-passphrase luksFormat "$RootPart"
	echo
	echo "Mounting new LUKS volume to continue preparation. Please unlock below:"
	${cmd} cryptsetup luksOpen "$RootPart" luksvol
}

function prepare_format() {
	if [[ "$DEBUG" ]]; then
		local cmd="echo"
	else
		local cmd=""
	fi

	echo "Formatting EFI: $EFIPart"
	${cmd} mkfs.fat -F32 -n "EFI" "$EFIPart"
  
  if [[ "$BootPart" != "@boot" ]]; then
  	echo "Formatting Boot: $BootPart"
  	${cmd} mkfs.ext4 -FF -L "Boot" "$BootPart"
  fi

  if [[ "$ENCRYPTION" ]]; then
  	echo "Formatting Root: /dev/mapper/luksvol"
  	${cmd} mkfs.btrfs -f -L "System" /dev/mapper/luksvol
  else
  	echo "Formatting Root: $RootPart"
  	${cmd} mkfs.btrfs -f -L "System" "$RootPart"
  fi
}

function create_subvolumes() {
	local -a subvols=("@")

	if [[ "$DEBUG" ]]; then
		local cmd="echo"
	else
		local cmd=""
	fi

	if [[ "$BootPart" == "@boot" ]]; then
		subvols+=("@boot")
	fi

    subvols+=("@swap")
	subvols+=("${!SUBVOLS_DEFAULT[@]}")

    [[ ! -d /mnt/btrfs ]] && ${cmd} mkdir -p /mnt/btrfs

	if [[ "$ENCRYPTION" ]]; then
		${cmd} mount "/dev/mapper/luksvol" /mnt/btrfs
	else
		${cmd} mount "$RootPart" /mnt/btrfs
	fi

	for subvol in "${subvols[@]}"
	do
		${cmd} btrfs subvolume create /mnt/btrfs/"$subvol"
	done

	${cmd} umount /mnt/btrfs
	${cmd} rmdir /mnt/btrfs
}

function get_hibernate_size() {
	free --giga | awk '/^Mem:/{print $2}'
}

function prepare_target() {
	local subvol
	local rootmount
    local compopt

	if [[ "$DEBUG" ]]; then
		local cmd="echo"
	else
		local cmd=""
	fi

	[[ ! -d /mnt/gentoo ]] && ${cmd} mkdir /mnt/gentoo

	if [[ "$ENCRYPTION" ]]; then
		rootmount="/dev/mapper/luksvol"
		if [[ ! -b "$rootmount" ]]; then
			${cmd} cryptsetup luksOpen "$RootPart" luksvol
		fi
	else
		rootmount="$RootPart"
	fi

    if [[ "$COMPRESSION" ]]; then
        compopt=",compress=zstd:3"
	fi

	${cmd} mount -o noatime,space_cache=v2${compopt},ssd,subvol=@ "$rootmount" /mnt/gentoo

	for subvol in "${!SUBVOLS_DEFAULT[@]}"
	do
		${cmd} mkdir -p /mnt/gentoo/"${SUBVOLS_DEFAULT[$subvol]}"
	done

	${cmd} mkdir -p /mnt/gentoo/boot
	if [[ "$BootPart" == "@boot" ]]; then
		${cmd} mount -o noatime,space_cache=v2,ssd,subvol=@boot "$rootmount" /mnt/gentoo/boot
	else
		${cmd} mount "$BootPart" /mnt/gentoo/boot
	fi

	for subvol in "${!SUBVOLS_DEFAULT[@]}"
	do
		${cmd} mount -o noatime,space_cache=v2,ssd,subvol="$subvol" "$rootmount" /mnt/gentoo/"${SUBVOLS_DEFAULT[$subvol]}"
	done

	${cmd} mkdir -p /mnt/gentoo/efi
	${cmd} mount "$EFIPart" /mnt/gentoo/efi

    ${cmd} mkdir -p /mnt/gentoo/swap
	${cmd} mount -o noatime,ssd,subvol=@swap "$rootmount" /mnt/gentoo/swap

	if [[ "$SWAP" && ! -f /mnt/gentoo/swap/hibernate.swp ]]; then
		#${cmd} mkdir -p /mnt/gentoo/swap
		#${cmd} mount -o noatime,ssd,subvol=@swap "$rootmount" /mnt/gentoo/swap
		${cmd} btrfs filesystem mkswapfile --size "$(get_hibernate_size)g" --uuid clear /mnt/gentoo/swap/hibernate.swp
	fi
}

function stage_step() {
	local s4file SwapUUID SwapOffset
	local luksUUID rootUUID rootmount
	local cmd rd ramdisk

	if [[ "$DEBUG" ]]; then
		cmd="echo"
	else
		cmd=""
	fi

	#FIXME check $INSTALL_STAGE for http or local file
	if [[ "$INSTALL_STAGE" == http* ]]; then
		s4file="${INSTALL_STAGE##*/}"
		${cmd} curl -o "/mnt/gentoo/${s4file}" "$INSTALL_STAGE" || errormsg 1 "Failed to download stage4 archive"
		s4file="/mnt/gentoo/${INSTALL_STAGE##*/}"
	else
		s4file="${INSTALL_STAGE}"
	fi

	${cmd} git clone "https://github.com/erenfro/gen2stage4" "$HOME/gen2stage4" || errormsg 1 "Failed to download gen2stage4"
	${cmd} "$HOME/gen2stage4/gen2extract" -t /mnt/gentoo "$s4file"
	${cmd} rm -rf "$HOME/gen2stage4"

	#FIXME post stage extraction:

	if [[ "$DEBUG" ]]; then
		echo "genfstab -U /mnt/gentoo > /mnt/gentoo/etc/fstab"
	else
		genfstab -U /mnt/gentoo > /mnt/gentoo/etc/fstab
	fi

	if [[ "$ENCRYPTION" ]]; then
		eval "$(blkid -p --output export "$RootPart" | grep UUID)"
		luksUUID="$UUID"
		rootmount="/dev/mapper/luksvol"
		eval "$(blkid -p --output export "$rootmount" | grep UUID)"
		rootUUID="$UUID"

		${cmd} mkdir /mnt/gentoo/etc/dracut.d

		if [[ "$DEBUG" ]]; then
			echo "echo \"luksvol UUID=\"$luksUUID\" none luks\" >> /mnt/gentoo/etc/crypttab"
			if [[ ! -d "/mnt/gentoo/etc/dracut.conf.d" ]]; then
				echo "mkdir /mnt/gentoo/etc/dracut.conf.d &>/dev/null"
				echo "echo \"add_dracutmodules+=\\\" crypt dm rootfs-block \\\" > /mnt/gentoo/etc/dracut.conf.d/luks.conf"
				echo "echo \"kernel_cmdline+=\\\" root=UUID=$rootUUID rd.luks.uuid=$luksUUID \\\" >> /mnt/gentoo/etc/dracut.conf.d/luks.conf"
			fi
		else
			echo "luksvol UUID=$luksUUID none luks" >> /mnt/gentoo/etc/crypttab
			if [[ ! -d "/mnt/gentoo/etc/dracut.conf.d" ]]; then
				mkdir /mnt/gentoo/etc/dracut.d &>/dev/null
				echo "add_dracutmodules+=\" crypt dm rootfs-block \"" > /mnt/gentoo/etc/dracut.conf.d/luks.conf
				echo "kernel_cmdline+=\" root=UUID=$rootUUID rd.luks.uuid=$luksUUID " >> /mnt/gentoo/etc/dracut.conf.d/luks.conf
			fi
		fi
	fi

	if [[ "$SWAP" ]]; then
		if [[ "$DEBUG" ]]; then
			echo "echo \"/swap/hibernate.swp none swap defaults 0 0\" >> /mnt/gentoo/etc/fstab"
		else
			echo "/swap/hibernate.swp none swap defaults 0 0" >> /mnt/gentoo/etc/fstab
		fi
		SwapUUID=$(grep btrfs /mnt/gentoo/etc/fstab | head -n1 | cut -f1)
		SwapOffset=$(btrfs inspect-internal map-swapfile -r /mnt/gentoo/swap/hibernate.swp)
        ${cmd} sed -i "s/^#GRUB_CMDLINE_LINUX_DEFAULT=.*/GRUB_CMDLINE_LINUX_DEFAULT=\"resume=${SwapUUID} resume_offset=${SwapOffset}\"/g" /mnt/gentoo/etc/default/grub
	fi

	if [[ -f "/mnt/gentoo/etc/machine-id" ]]; then
		${cmd} rm -f /mnt/gentoo/etc/machine-id
	fi
	${cmd} arch-chroot /mnt/gentoo systemd-machine-id-setup
	${cmd} arch-chroot /mnt/gentoo emerge --sync
	${cmd} arch-chroot /mnt/gentoo grub-install --efi-directory=/efi
	#arch-chroot /mnt/gentoo dracut --host-only
	while read -r rd; do
		rd="$(basename "$rd")"
		${cmd} arch-chroot --force "/boot/initramfs-${rd}.img" --kver "$rd"
	done < <(find /mnt/gentoo/lib/modules -mindepth 1 -maxdepth 1 -type d)
	${cmd} arch-chroot /mnt/gentoo grub-mkconfig -o /boot/grub/grub.cfg
}

function show_options() {
    echo "System Disk:    $RootDisk"
	echo "Root Partition: $RootPart"
	echo "Boot Partition: $BootPart"
	echo "EFI Partition:  $EFIPart"

    if [[ "$INSTALL_MODE" == "stage" ]]; then
        echo "Install Mode:   Stage"
        echo "STAGE File:     $INSTALL_STAGE"
    elif [[ "$INSTALL_MODE" == "clean" ]]; then
        echo "Install Mode:   Clean"
    else
        echo "Install Mode:   Normal"
    fi

	if [[ "$ENCRYPTION" ]]; then
		echo "Encryption:     Enabled"
	else
		echo "Encryption:     Disabled"
	fi

    if [[ "$COMPRESSION" ]]; then
		echo "Compression:    Enabled"
	else
		echo "Compression:    Disabled"
	fi

	if [[ "$SWAP" ]]; then
		echo "Swap:           Enabled"
	else
		echo "Swap:           Disabled"
	fi

	if [[ "$DEBUG" ]]; then
		echo "Debug:          Enabled"
	else
		echo "Debug:          Disabled"
	fi
	echo
}

function install_prep() {
	show_options

	read -rsn1 -p"Preparation: To proceed, press enter to continue." proceed
	echo
	if [[ "$proceed" != "" ]]; then
		echo "Aborting."
		exit 1
	fi

    echo "Creating Partitions on ${RootDisk}..."
    prepare_disk

	if [[ "$ENCRYPTION" ]]; then
		echo
		echo "Initiallizing LUKSv2 on ${RootPart}:"
		create_luks
	fi

	echo
	echo "Formatting Filesystems..."
	prepare_format

	echo
	echo "Preparing Subvolumes..."
	create_subvolumes

	echo
	echo "Preparing Installation Target..."
	prepare_target

    echo
    echo "Installation is ready and mounted in /mnt/gentoo for stage3 or stage4 install"
    echo "Please verify all went well, and re-run this script with --stage <stagefile>"
    echo "as appropriate, to continue."
}

function install_mount() {
	show_options

	read -rsn1 -p"Mounting: To proceed, press enter to continue." proceed
	echo
	if [[ "$proceed" != "" ]]; then
		echo "Aborting."
		exit 1
	fi

    echo "Mounting Partitions on ${RootDisk}..."
	prepare_target
}

function install_umount() {
	show_options

	read -rsn1 -p"Unounting: To proceed, press enter to continue." proceed
	echo
	if [[ "$proceed" != "" ]]; then
		echo "Aborting."
		exit 1
	fi

    echo "Mounting Partitions on ${RootDisk}..."
	umount -R /mnt/gentoo
}

function install_stage() {
	show_options

	read -rsn1 -p"Stage-Installation: To proceed, press enter to continue." proceed
	echo

	if [[ "$proceed" != "" ]]; then
		echo "Aborting."
		exit 1
	fi

	echo
	echo "Running Stage-Mode Installation Steps..."
    stage_step
}

function install_cleanup() {
	show_options

    echo "!!!WARNING!!! This is about to destructively wipe ${RootDisk}!"
	read -r -p"Enter YES to continue: " proceed
	echo
	if [[ "${proceed,,}" != "yes" ]]; then
		echo "Aborting."
		exit 1
	fi

	if [[ "$DEBUG" ]]; then
		local cmd="echo"
	else
		local cmd=""
	fi

	for subvol in "${!SUBVOLS_DEFAULT[@]}"
	do
		${cmd} umount /mnt/gentoo/"${SUBVOLS_DEFAULT[$subvol]}"
	done

	${cmd} umount /mnt/gentoo/efi
	${cmd} umount /mnt/gentoo/boot
	${cmd} umount /mnt/gentoo/swap
	${cmd} umount /mnt/gentoo

	if [[ "$ENCRYPTION" ]]; then
		${cmd} cryptsetup luksClose luksvol
	fi

    ${cmd} dd if=/dev/zero of="${RootDisk}" bs=1024 count=10
}


declare -a POSITIONAL_ARGS=()
declare INSTALL_MODE="normal"

while [[ $# -gt 0 ]]; do
	case $1 in
        -c|--compression)
            COMPRESSION=true
            shift
            ;;
		-d|--debug)
			DEBUG=true
			shift
			;;
		-e|--encryption)
			ENCRYPTION=true
			shift
			;;
		-h|--help)
			show_help
			;;
		-s|--swap)
			SWAP=true
			shift
			;;
		-u|--uki)
			UKI=true
			shift
			;;
		--mount)
			INSTALL_MODE=mount
			shift
			;;
		--umount)
			INSTALL_MODE=umount
			shift
			;;
		--stage)
			INSTALL_MODE=stage
            INSTALL_STAGE=$2
			shift 2
			;;
		--clean)
			INSTALL_MODE=clean
			shift
			;;
		*)
			POSITIONAL_ARGS+=("$1") # save positional arg
			shift
			;;
	esac
done

set -- "${POSITIONAL_ARGS[@]}" # restore positional parameters

RootDisk="$1"

if [[ ! -b "$RootDisk" ]]; then
    echo "ERROR: Invalid parameters. See --help for help"
    exit 3
else
    if [[ "$ENCRYPTION" ]]; then
        diskdev="${RootDisk##*/}"
        if [[ "$diskdev" =~ ^nvme.* ]]; then
            RootPart="${RootDisk}p3"
            BootPart="${RootDisk}p2"
            EFIPart="${RootDisk}p1"
        else
            RootPart="${RootDisk}3"
            BootPart="${RootDisk}2"
            EFIPart="${RootDisk}1"
        fi
    else
        BootPart="@boot"
        if [[ "$diskdev" =~ ^nvme.* ]]; then
            RootPart="${RootDisk}p2"
            EFIPart="${RootDisk}p1"
        else
            RootPart="${RootDisk}2"
            EFIPart="${RootDisk}1"
        fi
    fi

	case "$INSTALL_MODE" in
		normal) install_prep;;
		mount)	install_mount;;
		umount)	install_umount;;
		stage)	install_stage;;
		clean)  install_cleanup;;
		*)
			echo "Error, unknown installation mode detected."
			exit 3
			;;
	esac
fi