#! /usr/bin/env nix-shell
#! nix-shell -i bash -p bash parted cryptsetup btrfs-progs
# shellcheck shell=bash

set -euo pipefail

read -p "enter the name of the storage device to partition (e.g. /dev/nvme0n1): " -r
drive=$REPLY

if [[ $drive =~ "nvme" ]]; then
  boot="${drive}p1"
  system="${drive}p2"
else
  boot="${drive}1"
  system="${drive}2"
fi

echo
echo "The following layout will be created:"
echo "
$drive
   |
   ├── $boot (boot partition, fat32)
   |
   └── $system (system partition, LUKS2 + btrfs)
           |
           ├── @ (mounted as /)
           |       |
           |       ├── /home (mounted @home subvolume)
           |       |
           |       ├── /boot (mounted boot partition)
           |       |
           |       └── /.snapshots (mounted @snapshots subvolume)
           |
           ├── @home (mounted as /home)
           |
           └── @snapshots (mounted as /.snapshots)
"

echo
echo "WARNING! Continuing will cause $drive to be formatted."
read -p "Do you really want to continue? [Y/n]" -n 1 -r
if [[ ! $REPLY =~ ^[Yy]$ ]]; then exit 1; fi

# create partition table
parted "$drive" -- mklabel gpt

# create boot partition
parted "$drive" -- mkpart ESP fat32 1MiB 512MiB
parted "$drive" -- set 1 esp on
mkfs.fat -F 32 -n boot "$boot"

# create system partition
parted "$drive" -- mkpart primary 512MiB 100%

# setup LUKS partition
cryptsetup -y -v luksFormat "$system"
cryptsetup open "$system" enc

# setup btrfs partition
mkfs.btrfs /dev/mapper/enc

# setup subvolumes
mount -t btrfs /dev/mapper/enc /mnt
btrfs subvolume create /mnt/@
btrfs subvolume create /mnt/@home
btrfs subvolume create /mnt/@snapshots
umount /mnt

# mount subvolumes
read -p "enter zstd compression level (default: 1): " -r
compression_level=${REPLY:-1}
mount -o subvol=@,compress-force=zstd:"$compression_level",noatime /dev/mapper/enc /mnt
mkdir /mnt/{home,.snapshots,boot}
mount -o subvol=@home,compress-force=zstd:"$compression_level",noatime /dev/mapper/enc /mnt/home
mount -o subvol=@snapshots,compress-force=zstd:"$compression_level",noatime /dev/mapper/enc /mnt/.snapshots
mount "$boot" /mnt/boot

# generate hardware-configuration.nix
nixos-generate-config --root /mnt

echo
echo "Partitions have been created and hardware-configuration.nix has been generated."
echo "WARNING: Some hardware-configuration.nix options might need to be set manually:"
echo "- add compress-force & noatime options"