#!/usr/env/bin bash


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

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

	if [[ -z "${swap_size}" ]]; then
		swap_size="$(grep MemTotal /proc/meminfo | awk '{print int(sqrt($2 / 1024 / 1024) * 2)}')"
	fi
	local swap_offset="$((swap_size + 1))"

	local label_crypt_luks="NixOS-Crypt"
	local label_swap="NixOS-Swap"
	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 >/dev/null 2>&1 enc || true
	dd if=/dev/zero of="${disk}" bs=512 count=1024 || 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}"
	# Swap Partition
	parted -a optimal "${disk}" -- mkpart "${label_swap}" linux-swap 1Gib "${swap_offset}GB"
	mkswap -L "${label_swap}" "${swap_partition}"
	swapon "${swap_partition}"
	# Nix Partition, where the OS will reside with our data
	parted -a optimal "${disk}" -- mkpart "${label_primary}" "${swap_offset}GiB" 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 unerypted 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

	mkdir -p /mnt/nix/persist
	# Finally, actually install NixOS
	nixos-install --flake "git+file:.#orion"
}

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

		    -s <int> | --swap <int> [OPTIONAL]
		    The size of the swap partition in gigabytes.
		    Default: sqrt(Current system ram in GB) * 2
		    Example:
		        $(basename "${0}") -s 32
	__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 encrypt_disk="no"
	local swap_size=""

	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"
			;;
		-s | --swap)
			shift
			if [[ "${1}" =~ ^[0-9]+$ ]]; then
				swap_size="${1}"
				if ((swap_size <= 0)); then
					printf "Invalid value passed for swap! Expected a non-zero positive number, got: %s\n" "${swap_size}"
					exit 1
				fi
				printf "Set swap size to '%s'\n" "${swap_size}"
			else
				printf "Invalid value passed for swap! Expected a number got: %s\n" "${1}"
				exit 1
			fi
			;;
		-?*)
			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
	read -r -s -n 1 -p "Press any key to begin NixOS installation to '${disk}'!"
	printf "Installing in "
	for i in {3..1}; do
		printf "%s.." "${i}"
		sleep 1
	done
	printf "\n"
	install "${disk}" "${encrypt_disk}" "${swap_size}"
}

main "${@}"