#!/bin/bash

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

declare -rA SUBVOLS_DEFAULT=(
  ["@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 "-s, --swap        Enable Swap/Hibernation support."
	echo "-d, --debug       Enable DEBUG mode for testing."
	exit 0
}

function create_subvolumes() {
	local -a subvols

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

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

	${cmd} mount "$RootPart" /mnt

	if [[ "$BootPart" == "@boot" ]]; then
        ${cmd} btrfs subvolume create /mnt/@boot
        ${cmd} rsync -avhHi --delete-after "/mnt/@/boot/" "/mnt/@boot/"
	fi

	if [[ "$SWAP" ]]; then
		${cmd} btrfs subvolume create /mnt/@swap
	fi

	for subvol in "${subvols[@]}"
	do
		${cmd} btrfs subvolume create /mnt/"$subvol"
        if [[ -d "/mnt/@/${SUBVOLS_DEFAULT[$subvol]}" || "DEBUG" ]]; then
            ${cmd} rsync -avhHi --delete-after "/mnt/@/${SUBVOLS_DEFAULT[$subvol]}/" "/mnt/$subvol/"
        fi
	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

	rootmount="$RootPart"

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

	for subvol in "${!SUBVOLS_DEFAULT[@]}"
	do
        if [[ ! -d "/target/${SUBVOLS_DEFAULT[$subvol]}" || "$DEBUG" ]]; then
		    ${cmd} mkdir -p /target/"${SUBVOLS_DEFAULT[$subvol]}"
        fi
	done

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

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

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

    if [[ ! -d "/target/boot/efi" ]]; then
	    ${cmd} mkdir -p /target/boot/efi
    fi
	${cmd} mount "$EFIPart" /target/boot/efi

	if [[ "$SWAP" ]]; then
        if [[ ! -d "/target/swap" ]]; then
		    ${cmd} mkdir -p /target/swap
        fi
		${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/fstab"
        echo "==== genfstab result : begin ===="
        genfstab -U /target
        echo "==== genfstab result : end ===="
        echo
	else
		genfstab -U /target > /target/etc/fstab
	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=${SwapUUID} resume_offset=${SwapOffset}&/" /target/etc/default/grub
	fi
}

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

	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

	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
}

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

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

	if [[ "$SWAP" ]]; then
		${cmd} umount /target/swap
	fi

	${cmd} umount /target/boot/efi
	${cmd} umount /target/boot
	${cmd} umount /target

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

	${cmd} dd if=/dev/zero of="${RootPart}" bs=1024 count=10
	${cmd} dd if=/dev/zero of="${EFIPart}" bs=1024 count=10
	if [[ "$BootPart" != "@boot" ]]; then
		${cmd} dd if=/dev/zero of="${BootPart}" bs=1024 count=10
	fi
}


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

while [[ $# -gt 0 ]]; do
	case $1 in
		-h|--help)
			show_help
			;;
		-s|--swap)
			SWAP=true
			shift
			;;
		-d|--debug)
			DEBUG=true
			shift
			;;
		--expert)
			INSTALL_MODE=expert
			shift
			;;
		--clean)
			INSTALL_MODE=clean
			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;;
		clean)  install_cleanup;;
		*)
			echo "Error, unknown installation mode detected."
			exit 3
			;;
	esac
fi