#!/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 "--clean           Cleanup disk for clean slate"
	exit 0
}

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
            parted "$RootDisk" --script --align optimal -- "$RootDisk" \
                                        mklabel gpt \
                                        mkpart primary 1MiB 100MiB \
                                        mkpart primary 100MiB 2048MiB \
                                        mkpart primary 2148MiB -2048s \
                                        set 1 esp
            #mkfs.vfat -F 32 -n "EFI" "$RootPart"
            #mkfs.ext4 -L "Boot" -m 0 "$EFIPart"
            #mkfs.btrfs -L "System" "$RootPart"
        else
            parted "$RootDisk" --script --align optimal -- "$RootDisk" \
                                        mklabel gpt \
                                        mkpart primary 1MiB 100MiB \
                                        mkpart primary 100MiB -2048s \
                                        set 1 esp
            #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 --cipher aes-xts-plain64 --hash sha512 --use-random --verify-passphrase luksFormat "$RootPart"
	${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 -L "System" /dev/mapper/luksvol
  else
  	echo "Formatting Root: $RootPart"
  	${cmd} mkfs.btrfs -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[@]}")

    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
}

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

	${cmd} mkdir /mnt/gentoo

	if [[ "$ENCRYPTION" ]]; then
		rootmount="/dev/mapper/luksvol"
	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" /tamnt/gentooget/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" ]]; 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 UUID PART_ENTRY_UUID
	local SwapUUID SwapOffset

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

	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)"
		if [[ "$DEBUG" ]]; then
			echo "echo \"luksvol UUID=$UUID none luks\" >> /mnt/gentoo/etc/crypttab"
		else
			echo "luksvol UUID=$UUID none luks" >> /mnt/gentoo/etc/crypttab
		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
}

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_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 -rs -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
		-h|--help)
			show_help
			;;
		-e|--encryption)
			ENCRYPTION=true
			shift
			;;
        -c|--compression)
            COMPRESSION=true
            shift
            ;;
		-s|--swap)
			SWAP=true
			shift
			;;
		-d|--debug)
			DEBUG=true
			shift
			;;
		--stage)
			INSTALL_MODE=stage
            INSTALL_STAGE=$1
			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;;
		stage)	install_stage;;
		clean)  install_cleanup;;
		*)
			echo "Error, unknown installation mode detected."
			exit 3
			;;
	esac
fi