#!/usr/bin/env bash
# Script to setup a drive for a brand new NixOS installation.
# IMPORTANT: Before running this script, you'll need to create a LUKS partition using:
#	cryptsetup --label=nixos-crypt --type=luks2 luksFormat [partition]

set -e

if [ "$(id -u)" != "0" ]; then
   echo "This script must be run as root" 1>&2
   exit 1
fi

# Configuration parameters
ask_root_password=false	# Whether to prompt for a root user password
boot_partition="" # The drive to install the bootloader to
luks_partition="" # The drive partition to create the LUKS container on
root_partition="/dev/mapper/nixos-crypt" # The partition to install NixOS to

function usage() {
	echo "Usage: format-drives.sh [--boot boot-partition-path] [--luks luks-partition-path] [--ask-root-password]"
	echo "Options:"
	echo "	-h | --help	Show this help screen."
	echo "	-b | --boot <path>	      The path to the boot drive (e.g. /dev/nvme0n1p1)."
	echo "  -l | --luks <path>        The path to the partition to create the LUKS container on (e.g. /dev/nvme0n1p2)."
	echo "  -a | --ask-root-password  Sets a password for the root user account."
	exit 2
}

# Argument processing logic shamelessly stolen from https://stackoverflow.com/questions/192249/how-do-i-parse-command-line-arguments-in-bash
POSITIONAL_ARGS=()
while [[ $# -gt 0 ]]; do
  case "$1" in
	--ask-root-password|-a)
		ask_root_password=true
		shift
		;;
	--boot|-b)
		boot_partition="$2"
		shift
		shift
		;;
	--luks|-l)
		luks_partition="$2"
		shift
		shift
		;;
	--help|-h)
		usage
		shift
		;;
	*)
		POSITIONAL_ARGS+=("$1") # save positional arg
      	shift # past argument
		;;
	esac
done

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

cryptsetup luksOpen $luks_partition nixos-crypt
mkfs.btrfs -L nixos $root_partition
mount /dev/mapper/nixos-crypt /mnt
btrfs subvolume create /mnt/@
btrfs subvolume create /mnt/@home
btrfs subvolume create /mnt/@log
btrfs subvolume create /mnt/@nix
btrfs subvolume create /mnt/@swap
umount /mnt

mount -o subvol=@ $root_partition /mnt
mkdir -p /mnt/{boot,home,var/log,nix,swap}
mount $boot_partition /mnt/boot
mount -o subvol=@home $root_partition /mnt/home
mount -o subvol=@log $root_partition /mnt/var/log
mount -o subvol=@nix $root_partition /mnt/nix
mount -o subvol=@swap $root_partition /mnt/swap
echo "Disks partitioned and mounted to /mnt."

# Generate hardware-configuration.nix
nixos-generate-config --no-filesystems --dir /home/nixos
echo "Configuration files generated and saved to /home/nixos."

echo "Setup complete!"
echo "To install, set up your system's configuration files under ./hosts/yourHost and add it to flake.nix."
echo "Then, run the following command:"
echo "nixos-install --verbose --root /mnt --flake [path-to-flake.nix] $( (( ask_root_password == false )) && echo "--no-root-password" )"

exit 0