# shellcheck shell=bash

usage() {
  echo "Usage: $0 flake -m install|repair -h host [-k key] [-p password_file] [-s] [-c] [-r]"
  echo
  echo "Options:"
  echo "  flake               Directory containing the flake.nix file."
  echo "  -m mode             Mode: 'install' or 'repair'."
  echo "  -h host             Host to configure."
  echo "  -k key              Key file to copy to user config."
  echo "  -s                  Enroll secure boot keys on current device."
  echo "  -c                  Copy configuration to target."
  echo "  -r                  Reboot after completion."
  exit 1
}

check_root() {
  if [[ "$EUID" -ne 0 ]]; then
    echo "Please run the script as root."
    exit 1
  fi
}

check_network() {
  if ! ping -c 1 google.com &>/dev/null; then
    echo "Connect to a network before proceeding."
    exit 1
  fi
}

check_flake() {
  if [[ ! -f "$flake/flake.nix" ]]; then
    echo "flake.nix not found in $flake."
    exit 1
  fi
}

check_host() {
  if ! nix flake show --allow-import-from-derivation --quiet --json "$flake" 2>/dev/null | jq -e ".nixosConfigurations[\"$host\"]" &>/dev/null; then
    echo "Host '$host' not found in flake."
    exit 1
  fi
}

check_key() {
  if [[ -n "$key" ]] && [[ ! -f "$flake/submodules/secrets/domains/$key/key.txt" ]]; then
    echo "Key '$key' not found."
    exit 1
  fi
}

set_password_file() {
  SOPS_AGE_KEY_FILE="$flake/submodules/secrets/domains/$key/key.txt"
  export SOPS_AGE_KEY_FILE
  install -m 600 /dev/null /tmp/keyfile
  sops --decrypt --extract "['luks']" "$flake/submodules/secrets/hosts/$host/secrets.yaml" > /tmp/keyfile
  unset SOPS_AGE_KEY_FILE
}

prepare_disk() {
  local disko_mode="$1"
  mkdir -p /mnt
  root=$(mktemp -d /mnt/install.XXXXXX)
  disko -m "$disko_mode" --yes-wipe-all-disks --root-mountpoint "$root" "$flake/hosts/$host/format.nix"
}

copy_sops_keys() {
  mkdir -p "$root/persist/state/etc/ssh"
  cp -f "$flake/submodules/secrets/hosts/$host/ssh_host_ed25519_key" "$root/persist/state/etc/ssh/ssh_host_ed25519_key"

  for path in "$flake/hosts/$host/users"/*; do
    if [[ -z "$key" ]]; then
      continue
    fi

    local user
    user=$(basename "$path")

    mkdir -p "$root/persist/state/home/$user/.config/sops-nix"
    cp -f "$flake/submodules/secrets/domains/$key/key.txt" "$root/persist/state/home/$user/.config/sops-nix/key.txt"

    owner=$(cat "$flake/hosts/$host/users/$user/uid")
    group=100
    chown "$owner:$group" \
      "$root/persist/state/home/$user" \
      "$root/persist/state/home/$user/.config" \
      "$root/persist/state/home/$user/.config/sops-nix" \
      "$root/persist/state/home/$user/.config/sops-nix/key.txt"
  done
}

copy_secure_boot_keys() {
  mkdir -p "$root/persist/state/var/lib/sbctl/keys"/{db,KEK,PK}

  SOPS_AGE_KEY_FILE="$flake/submodules/secrets/domains/$key/key.txt"
  export SOPS_AGE_KEY_FILE

  sops --decrypt --extract "['guid']" "$flake/submodules/secrets/domains/lanzaboote/secrets.yaml" > "$root/persist/state/var/lib/sbctl/GUID"
  sops --decrypt --extract "['keys']['kek']['key']" "$flake/submodules/secrets/domains/lanzaboote/secrets.yaml" > "$root/persist/state/var/lib/sbctl/keys/KEK/KEK.key"
  sops --decrypt --extract "['keys']['kek']['pem']" "$flake/submodules/secrets/domains/lanzaboote/secrets.yaml" > "$root/persist/state/var/lib/sbctl/keys/KEK/KEK.pem"
  sops --decrypt --extract "['keys']['pk']['key']" "$flake/submodules/secrets/domains/lanzaboote/secrets.yaml" > "$root/persist/state/var/lib/sbctl/keys/PK/PK.key"
  sops --decrypt --extract "['keys']['pk']['pem']" "$flake/submodules/secrets/domains/lanzaboote/secrets.yaml" > "$root/persist/state/var/lib/sbctl/keys/PK/PK.pem"
  sops --decrypt --extract "['keys']['db']['key']" "$flake/submodules/secrets/domains/lanzaboote/secrets.yaml" > "$root/persist/state/var/lib/sbctl/keys/db/db.key"
  sops --decrypt --extract "['keys']['db']['pem']" "$flake/submodules/secrets/domains/lanzaboote/secrets.yaml" > "$root/persist/state/var/lib/sbctl/keys/db/db.pem"

  chmod 400 "$root/persist/state/var/lib/sbctl/keys"/*/*

  unset SOPS_AGE_KEY_FILE

  mkdir -p "$root/var/lib/sbctl"
  mount --bind -o X-fstrim.notrim,x-gvfs-hide "$root/persist/state/var/lib/sbctl" "$root/var/lib/sbctl"
}

install_nixos() {
  nixos-install --root "$root" --flake "$flake#$host" --no-root-passwd
}

enroll_secure_boot() {
  sbctl enroll-keys --microsoft
}

copy_config() {
  echo "Copying configuration..."
  mkdir -p "$root/persist/user/etc"
  rm -rf "$root/persist/user/etc/nixos"
  cp -r "$flake" "$root/persist/user/etc/nixos"
}

cleanup() {
  rm -f /tmp/keyfile
  if [[ -d "$root" ]]; then umount "$root/var/lib/sbctl"; fi
  if [[ -n "$host" ]]; then disko -m "unmount" "$flake/hosts/$host/format.nix"; fi
  if [[ -d "$root" ]]; then rmdir "$root"; fi
}

main() {
  check_root
  check_network

  if [[ "$#" -lt 1 ]]; then usage; fi

  flake="$(realpath "$1")"
  check_flake
  shift

  mode=""
  host=""
  key=""
  enroll_secure_boot_flag="false"
  copy_config_flag="false"
  reboot_flag="false"

  while getopts "m:h:k:scr" opt; do
    case "$opt" in
      m) mode="$OPTARG" ;;
      h) host="$OPTARG" ;;
      k) key="$OPTARG" ;;
      s) enroll_secure_boot_flag="true" ;;
      c) copy_config_flag="true" ;;
      r) reboot_flag="true" ;;
      *) usage ;;
    esac
  done

  if [[ -z "$mode" || -z "$host" ]]; then usage; fi

  check_host
  check_key
  set_password_file

  case "$mode" in
    install) prepare_disk "destroy,format,mount";;
    repair) prepare_disk "mount";;
    *)
      echo "Invalid mode: $mode"
      usage
      ;;
  esac

  copy_sops_keys
  copy_secure_boot_keys

  install_nixos

  [[ "$enroll_secure_boot_flag" == "true" ]] && enroll_secure_boot
  [[ "$copy_config_flag" == "true" ]] && copy_config

  cleanup

  [[ "$reboot_flag" == "true" ]] && reboot
}

main "$@"