#!/usr/bin/env bash
set -e

echo "========================================="
echo "   Interactive NixOS Flake Installer     "
echo "========================================="

# Prompt for the flake URL and GitHub PAT
read -p "Enter the target flake (e.g., github:izziekitty/nixos#izzie-nixos-god): " FLAKE
read -s -p "Enter your GitHub Personal Access Token (PAT): " GITHUB_PAT
echo ""

echo ""
echo "Available Disks:"
echo "----------------"
lsblk -d -o NAME,SIZE,MODEL,TYPE | grep -i disk
echo "----------------"

read -p "Enter the MAIN OS drive (e.g., nvme0n1): " OS_DISK
if [ ! -b "/dev/${OS_DISK}" ]; then echo "Invalid device."; exit 1; fi

read -p "Do you want /home on a separate drive? (y/n): " SEP_HOME
if [[ "${SEP_HOME}" == "y" ]]; then
  read -p "Enter the /home drive (e.g., nvme1n1): " HOME_DISK
else
  HOME_DISK="${OS_DISK}"
fi

read -p "Do you want /dec/nixos on a separate drive? (y/n): " SEP_DEC
if [[ "${SEP_DEC}" == "y" ]]; then
  read -p "Enter the /dec/nixos drive (e.g., nvme2n1 or same as /home): " DEC_DISK
else
  DEC_DISK="${OS_DISK}"
fi

echo ""
echo "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
read -p "CRITICAL WARNING: This will format the selected drives. Type YES to continue: " CONFIRM
if [ "${CONFIRM}" != "YES" ]; then echo "Aborted."; exit 0; fi

generate_disk() {
  local disk_name="$1"
  local device="$2"
  local is_os="$3"
  local has_home="$4"
  local has_dec="$5"
  local fs_label="$6"

  echo "    ${disk_name} = {"
  echo "      type = \"disk\";"
  echo "      device = \"/dev/${device}\";"
  echo "      content = {"
  echo "        type = \"gpt\";"
  echo "        partitions = {"

  if [ "${is_os}" == "true" ]; then
    echo "          ESP = {"
    echo "            size = \"1024M\";"
    echo "            type = \"EF00\";"
    echo "            content = {"
    echo "              type = \"filesystem\";"
    echo "              format = \"vfat\";"
    echo "              mountpoint = \"/boot\";"
    echo "              mountOptions = [ \"umask=0077\" ];"
    echo "            };"
    echo "          };"
  fi

  echo "          root = {"
  echo "            size = \"100%\";"
  echo "            content = {"
  echo "              type = \"btrfs\";"
  echo "              extraArgs = [ \"-f\" \"-L\" \"${fs_label}\" ];"
  echo "              subvolumes = {"

  if [ "${is_os}" == "true" ]; then
    echo "                \"root\" = { mountpoint = \"/\"; mountOptions = [ \"compress=zstd\" \"noatime\" ]; };"
    echo "                \"nix\" = { mountpoint = \"/nix\"; mountOptions = [ \"compress=zstd\" \"noatime\" ]; };"
    echo "                \"persist\" = { mountpoint = \"/persist\"; mountOptions = [ \"compress=zstd\" \"noatime\" ]; };"
  fi
  if [ "${has_home}" == "true" ]; then
    echo "                \"home\" = { mountpoint = \"/home\"; mountOptions = [ \"compress=zstd\" \"noatime\" ]; };"
  fi
  if [ "${has_dec}" == "true" ]; then
    echo "                \"nixos\" = { mountpoint = \"/dec/nixos\"; mountOptions = [ \"compress=zstd\" \"noatime\" ]; };"
  fi

  echo "              };"
  echo "            };"
  echo "          };"
  echo "        };"
  echo "      };"
  echo "    };"
}

echo "{" > /tmp/disko.nix
echo "  disko.devices.disk = {" >> /tmp/disko.nix

os_has_home="false"; os_has_dec="false"
if [ "${HOME_DISK}" == "${OS_DISK}" ]; then os_has_home="true"; fi
if [ "${DEC_DISK}" == "${OS_DISK}" ]; then os_has_dec="true"; fi
generate_disk "main" "${OS_DISK}" "true" "${os_has_home}" "${os_has_dec}" "NIXOS_MAIN" >> /tmp/disko.nix

if [ "${HOME_DISK}" != "${OS_DISK}" ]; then
  home_has_dec="false"
  if [ "${DEC_DISK}" == "${HOME_DISK}" ]; then home_has_dec="true"; fi
  generate_disk "home_drive" "${HOME_DISK}" "false" "true" "${home_has_dec}" "NIXOS_HOME" >> /tmp/disko.nix
fi

if [ "${DEC_DISK}" != "${OS_DISK}" ] && [ "${DEC_DISK}" != "${HOME_DISK}" ]; then
  generate_disk "dec_drive" "${DEC_DISK}" "false" "false" "true" "NIXOS_DEC" >> /tmp/disko.nix
fi

echo "  };" >> /tmp/disko.nix
echo "}" >> /tmp/disko.nix

echo "Running Disko partitioning..."
nix run --extra-experimental-features "nix-command flakes" github:nix-community/disko -- --mode disko /tmp/disko.nix

echo "Preparing impermanence snapshot..."
mkdir -p /tmp/btrfs-top
mount -t btrfs -L NIXOS_MAIN -o subvol=/ /tmp/btrfs-top
btrfs subvolume snapshot -r /tmp/btrfs-top/root /tmp/btrfs-top/root-blank
umount /tmp/btrfs-top
echo "Snapshot root-blank created successfully."

echo "Installing NixOS from ${FLAKE}..."
# The installation command now utilizes the GitHub PAT for private repo access
nixos-install --flake "${FLAKE}" --no-root-passwd --option access-tokens "github.com=${GITHUB_PAT}"

echo "Installation complete! You can now reboot."