#!/bin/bash

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

function show_help() {
	echo "Usage: $0 <root> <boot> <efi> [<options>...]"
	echo ""
	echo "Positional Arguments:"
	echo "<root>    Root Partition used for LUKS or BtrFS filesystem"
	echo "<boot>    Boot Partition or @boot for subvolume in BtrFS"
	echo "<efi>     ESP Partition for booting"
	echo ""
	echo "Options:"
	echo "-h, --help        Help on this tool."
	echo "-e, --encryption  Enable LUKS encryption."
	echo "-s, --swap        Enable Swap/Hibernation support."
	echo "-d, --debug       Enable DEBUG mode for testing."
	exit 0
}

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 root /dev/mapper/luksvol
  else
  	echo "Formatting Root: $RootPart"
  	${cmd} mkfs.btrfs -L root "$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

	if [[ "$SWAP" ]]; then
		subvols+=("@swap")
	fi

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

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

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

	${cmd} umount /mnt
}

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

function prepare_target() {
	local subvol
	local rootmount

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

	${cmd} mkdir /target

	if [[ "$ENCRYPTION" ]]; then
		rootmount="/dev/mapper/luksvol"
	else
		rootmount="$RootPart"
	fi

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

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

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

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

	${cmd} mkdir -p /target/boot/efi
	${cmd} mount "$EFIPart" /target/boot/efi

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

function expert_step() {
	local UUID PART_ENTRY_UUID
	local SwapUUID SwapOffset

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

	${cmd} apt install -y arch-install-scripts

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

	if [[ "$ENCRYPTION" ]]; then
		eval "$(blkid -p --output export "$RootPart" | grep UUID)"
		if [[ "$DEBUG" ]]; then
			echo "echo \"luksvol UUID=$UUID none luks\" >> /target/etc/crypttab"
		else
			echo "luksvol UUID=$UUID none luks" >> /target/etc/crypttab
		fi
	fi

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

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

	if [[ "$ENCRYPTION" ]]; then
		echo "Encryption:     Enabled"
	else
		echo "Encryption:     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_normal() {
	show_options

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

	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 "Ready for installation! Run a terminal and start the following:"
	echo "sudo live-installer-expert-mode"
	echo
	echo "Once it's at the expert mode step, re-run this command with --expert"
}

function install_expert() {
	show_options

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

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

	echo
	echo "Running Expert-Mode Installation Steps..."
	expert_step
}

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

while [[ $# -gt 0 ]]; do
	case $1 in
		-h|--help)
			show_help
			;;
		-e|--encryption)
			ENCRYPTION=true
			shift
			;;
		-s|--swap)
			SWAP=true
			shift
			;;
		-d|--debug)
			DEBUG=true
			shift
			;;
		--expert)
			INSTALL_MODE=expert
			shift
			;;
		*)
			POSITIONAL_ARGS+=("$1") # save positional arg
			shift
			;;
	esac
done

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

RootPart="$1"
BootPart="$2"
EFIPart="$3"

if [[ -z "$RootPart" || -z "$BootPart" || -z "$EFIPart" ]]; then
	echo "ERROR: Invalid parameters. See --help for help"
	exit 3
else
	if [[ "$ENCRYPTION" && "$BootPart" == "@boot" ]]; then
		echo "While encryption is enabled, using @boot subvolume will cause issues."
		echo "Please prepare and set boot volume for /boot while using encryption."
		exit 4
	fi

	case "$INSTALL_MODE" in
		normal) install_normal;;
		expert)	install_expert;;
		*)
			echo "Error, unknown installation mode detected."
			exit 3
			;;
	esac
fi