#!/usr/env/bin bash

install() {
	set -euo pipefail
	local disk="${1}"
	local encrypt_disk="${2}"
	local host="${3}"

	local boot_partition="${disk}1"
	local primary_partition="${disk}2"
	# The size is large because I'd like to be able to hibernate my laptop in its entirety. I have 64 GB of ram.

	local label_crypt_luks="NixOS-Crypt"
	local label_primary="NixOS-Primary"
	local label_boot="NixOS-Boot"

	swapoff -a >/dev/null 2>&1 || true
	umount /mnt/**/* >/dev/null 2>&1 || true
	umount /mnt/* >/dev/null 2>&1 || true
	umount /mnt >/dev/null 2>&1 || true
	cryptsetup close enc >/dev/null 2>&1 || true
	wipefs -a "${disk}" || true

	### Partition The Disk
	parted "${disk}" -- mklabel gpt
	# Boot partition
	parted -a optimal "${disk}" -- mkpart ESP fat32 1MiB 1GiB
	parted "${disk}" -- set 1 boot on
	mkfs.vfat "${boot_partition}"
	fatlabel "${boot_partition}" "${label_boot}"
	# Primary Partition
	parted -a optimal "${disk}" -- mkpart "${label_primary}" 1Gib 100%

	### Encrypt
	if [[ "${encrypt_disk}" == "yes" ]]; then
		cryptsetup --verify-passphrase -v luksFormat "${primary_partition}"
		cryptsetup config "${primary_partition}" --label "${label_crypt_luks}"
		# Have to decrypt it so we can actually get other things setup
		local crypt_open_name="enc"
		cryptsetup open "${primary_partition}" "${crypt_open_name}"
		primary_partition="/dev/mapper/${crypt_open_name}"
	fi

	### BTRFS Setup
	# Go ahead and make the unencrypted BTRFS
	mkfs.btrfs -f -L "${label_primary}" "${primary_partition}"

	# Mount it
	mount -t btrfs "${primary_partition}" /mnt

	# Create our subvolumes
	btrfs subvolume create "/mnt/@nix"
	umount /mnt

	### Final Mountings
	# Mount tmpfs to mnt
	mount -t tmpfs -o mode=755 none /mnt

	# Create our directories
	mkdir /mnt/{"boot","nix"}
	# Mount our boot partition
	mount -t vfat -o defaults,noatime "${boot_partition}" /mnt/boot

	# Mount our btrfs subvolumes individually with some btrfs options
	# NOTE: On high performance NVME SSDs with a beefy CPU it may be worth considering ZLO compression instead of ZSTD. In
	# many cases ZLO is more performant, especially when writing, than ZSTD while having a somewhat worse comrpession ratio.
	# WARN: ZLO *may* be a good solution, it can be VERY slow on incompressible data. Something to keep in mind.
	mount -t btrfs -o noatime,compress=zstd,subvol=@nix "${primary_partition}" /mnt/nix

	# Persistence dir, dirs not to be wiped on reboot stored here
	mkdir -p /mnt/nix/persist
	# Actually install NixOS
	nixos-install --flake "git+file:.#${host}"
	# Finally, clone the flake into the persistent nixos config
	git clone . /mnt/nix/persist/etc/nixos/
}

usage() {
	cat <<-__EOF__
		Usage: $(basename "${0}") -d "/dev/DISK_HERE" -s 32 -e
		    -h | -? | --help
		    Shows this menu.
		    Example:
		        $(basename "${0}") -h

		    -d <string> | --disk <string> [REQUIRED]
		    Path to a disk to install NixOS to.
		    Default: none
		    Example:
		        $(basename "${0}") -d /dev/vda

		    -e | --encrypt [OPTIONAL]
		    Enable disk encryption during install. Note, this will prompt you for the encryption password.
		    Default: disabled
		    Example:
		        $(basename "${0}") -e
	__EOF__
}

running_as_root() {
	((EUID == 0))
}

main() {
	if ! running_as_root; then
		printf "This script MUST be ran as root, exiting!\n"
		exit 1
	fi

	local disk=""
	local host=""
	local encrypt_disk="no"

	while :; do
		case ${1} in
		-h | -\? | --help)
			usage # Display a usage synopsis.
			exit
			;;
		--) # End of all options.
			break
			;;
		-d | --disk)
			shift
			disk="${1}"
			if [[ ! -b "${disk}" ]]; then
				printf "Unable to read disk '%s', check that it exists, is a block device, and that you have write and read permissions for it!" "${disk}"
				exit 1
			fi
			printf "Setting installation disk to '%s'\n" "${disk}"
			;;
		-e | --encrypt)
			encrypt_disk="yes"
			printf "Enabled disk encryption\n"
			;;
		-H | --host)
			shift
			host="${1}"
			printf "Set host to '%s'\n" "${host}"
			;;
		-?*)
			printf 'Unknown option: %s\n' "$1" >&2
			usage
			exit 1
			;;
		*) # Default case: No more options, so break out of the loop.
			break ;;
		esac
		shift
	done

	set -euo pipefail
	if [[ -z "${disk}" ]]; then
		printf "Value for 'disk' (-d) was not provided! See '--help' for usage!\n"
		exit 1
	fi
	if [[ -z "${host}" ]]; then
		printf "Value to 'host' (-H) was not provided! See '--help' for usage!\n"
		exit 1
	fi
	read -r -s -n 1 -p "Press any key to begin NixOS installation to '${disk}' for host '${host}'!"
	# printf "\nInstalling in "
	# for i in {3..1}; do
	# 	printf "%s.." "${i}"
	# 	sleep 1
	# done
	# printf "\n"
	install "${disk}" "${encrypt_disk}" "${host}"
}

main "${@}"