usage() {
  echo "Usage: $0 flake -m install|repair -h host [-k key] [-p password_file] [-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 "  -p password_file    LUKS password file to use for encryption."
  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 --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}/secrets/${key}/key.txt" ]]; then
    echo "Key '${key}' not found."
    exit 1
  fi
}

set_password_file() {
  if [[ -n "${password_file}" ]]; then
    if [[ ! -f "${password_file}" ]]; then
      echo "LUKS key file '${password_file}' not found."
      exit 1
    fi

    ln -sf "${password_file}" /tmp/installer.key
  else
    echo "Enter password for LUKS encryption:"
    IFS= read -r -s password
    echo "Enter password again to confirm: "
    IFS= read -r -s password_check
    [ "${password}" != "${password_check}" ]
    echo -n "${password}" > /tmp/installer.key
    unset password password_check
  fi
}

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

copy_keys() {
  mkdir -p "${root}/persist/etc/ssh"
  cp "${flake}/hosts/${host}/secrets/ssh_host_ed25519_key" "${root}/persist/etc/ssh/ssh_host_ed25519_key"

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

    user=$(basename "${path}")
    mkdir -p "${root}/persist/home/${user}/.config/sops-nix"
    cp "${flake}/secrets/${key}/key.txt" "${root}/persist/home/${user}/.config/sops-nix/key.txt"
    uid=$(cat "${flake}/hosts/${host}/users/${user}/uid")
    gid=100
    chown -R "${uid}:${gid}" "${root}/persist/home/${user}"
  done
}

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

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

finish() {
  echo "Rebooting system..."
  trap - EXIT
  cleanup
  reboot
}

cleanup() {
  rm -f /tmp/installer.key
  if [[ -n "${host}" && -n "${device}" ]]; then disko -m "unmount" "${flake}/hosts/${host}/format.nix" --arg device "\"${device}\""; fi
  if [[ -d "${root}" ]]; then rmdir "${root}"; fi
}

check_root
check_network

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

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

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

while getopts "m:h:k:p:cr" opt; do
  case "${opt}" in
    m) mode="${OPTARG}" ;;
    h) host="${OPTARG}" ;;
    k) key="${OPTARG}" ;;
    p) password_file="${OPTARG}" ;;
    c) copy_config_flag="true" ;;
    r) reboot_flag="true" ;;
    *) usage ;;
  esac
done

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

check_host
check_key
until set_password_file; do echo "Passwords did not match, please try again."; done

device=$(grep -oP '(?<=device = ")[^"]+' "${flake}/hosts/${host}/default.nix")

case "${mode}" in
  install)
    prepare_disk "destroy,format,mount"
    copy_keys
    install
    if [[ "${copy_config_flag}" == "true" ]]; then copy_config; fi
    if [[ "${reboot_flag}" == "true" ]]; then finish; fi
    ;;
  repair)
    prepare_disk "mount"
    install
    if [[ "${reboot_flag}" == "true" ]]; then finish; fi
    ;;
  *)
    echo "Invalid mode: ${mode}"
    usage
    ;;
esac