#!/usr/bin/env bash

set -euo pipefail

showUsage() {
  cat <<EOF
Usage: $0 [OPTIONS]
  --mode MODE                 Specify the mode of operation. Valid modes are: format, mount.
                              Format will format the disk before installing.
                              Mount will mount the disk before installing.
                              Mount is useful for updating an existing system without losing data.
  -f, --flake FLAKE_URI#ATTR  Use the specified flake to install the NixOS configuration.
  --disk NAME DEVICE          Map the specified disk name to the specified device path.
  --dry-run                   Print the commands that would be run, but do not run them.
  --show-trace                Show the stack trace on error.
  -h, --help                  Show this help message.
  --extra-files SOURCE DEST   Copy the specified file or directory from the host into the NixOS configuration.
  --option NAME VALUE         Pass the specified option to Nix.
  --write-efi-boot-entries    Write EFI boot entries to the NVRAM of the system for the installed system.
                              Specify this option if you plan to boot from this disk on the current machine,
                              but not if you plan to move the disk to another machine.
  --system-config JSON        Merges the specified JSON object into the NixOS configuration.
  --mount-point PATH          Specify the mount point for the NixOS installation.
                              Default: /mnt/disko-install-root
EOF
}

serialiaseArrayToNix() {
  local -n array=$1
  nixExpr="{ "
  # Iterate over the associative array to populate the Nix attrset string
  for key in "${!array[@]}"; do
    value=${array[$key]}
    nixExpr+="\"${key//\"/\\\"}\" = \"${value//\"/\\\"}\"; "
  done
  nixExpr+="}"

  echo "$nixExpr"
}

readonly libexec_dir="${0%/*}"

nix_args=(
  --extra-experimental-features 'nix-command flakes'
  "--option" "no-write-lock-file" "true"
)
dry_run=
extraSystemConfig="{}"
diskoAttr=diskoScript
writeEfiBootEntries=false
mountPoint=/mnt/disko-install-root
declare -A diskMappings
declare -A extraFiles


parseArgs() {
  [[ $# -eq 0 ]] && {
    showUsage
    exit 1
  }

  while [[ $# -gt 0 ]]; do
    case "$1" in
    -d | --debug)
      set -x
      ;;
    -f | --flake)
      flake=$2
      shift
      ;;
    -h | --help)
      showUsage
      exit 0
      ;;
    --dry-run)
      dry_run=y
      ;;
    --show-trace)
      nix_args+=("$1")
      ;;
    --write-efi-boot-entries)
      writeEfiBootEntries=true
      ;;
    --mode)
      if [[ $# -lt 2 ]]; then
        echo "Option $1 requires an argument" >&2
        exit 1
      fi
      case $2 in
      format)
        diskoAttr=diskoScript
        ;;
      mount)
        diskoAttr=mountScript
        ;;
      *)
        echo "Invalid mode: $2" >&2
        echo "Valid modes are: format, mount" >&2
        exit 1
        ;;
      esac
      shift
      ;;
    --system-config)
      if [[ $# -lt 2 ]]; then
        echo "Option $1 requires one JSON argument." >&2
        exit 1
      fi
      # shellcheck disable=SC2034
      extraSystemConfig="$2"
      shift
      ;;
    --extra-files)
      if [[ $# -lt 3 ]]; then
        echo "Option $1 requires two arguments: source, destination" >&2
        exit 1
      fi
      extraFiles[$2]=$3
      shift
      shift
      ;;
    --option)
      if [[ $# -lt 3 ]]; then
        echo "Option $1 requires two arguments: key, value" >&2
        exit 1
      fi
      nix_args+=(--option "$2" "$3")
      shift
      shift
      ;;
    --disk)
      if [[ $# -lt 3 ]]; then
        echo "Option $1 requires two arguments: disk_name, device_path" >&2
        exit 1
      fi
      # shellcheck disable=SC2034
      diskMappings[$2]=$3
      shift
      shift
      ;;
    --mount-point)
      if [[ $# -lt 2 ]]; then
        echo "Option $1 requires an argument" >&2
        exit 1
      fi
      mountPoint=$2
      shift
      ;;
    *)
      echo "Unknown option: $1" >&2
      showUsage
      exit 1
      ;;
    esac
    shift
  done
}

cleanupMountPoint() {
  if mountpoint -q "${mountPoint}"; then
    umount -R "${mountPoint}"
  fi
  rmdir "${mountPoint}"
}

nixBuild() {
  if command -v nom-build > /dev/null; then
    nom-build "$@"
  else
    nix-build "$@"
  fi
}

maybeRun () {
  if [[ -z ${dry_run-} ]]; then
    "$@"
  else
    echo "Would run: $*"
  fi
}

main() {
  parseArgs "$@"

  if [[ -z ${flake-} ]]; then
    echo "Please specify the flake-uri with --flake to use for installation." >&2
    exit 1
  fi

  # check if we are root
  if [[ "$EUID" -ne 0 ]] && [[ -z ${dry_run-} ]]; then
    echo "This script must be run as root" >&2
    exit 1
  fi

  if [[ $flake =~ ^(.*)\#([^\#\"]*)$ ]]; then
    flake="${BASH_REMATCH[1]}"
    flakeAttr="${BASH_REMATCH[2]}"
  fi

  if [[ -e "$flake" ]]; then
    flake=$(realpath "$flake")
  fi

  if [[ -z ${flakeAttr-} ]]; then
    echo "Please specify the name of the NixOS configuration to be installed, as a URI fragment in the flake-uri." >&2
    echo 'For example, to use the output nixosConfigurations.foo from the flake.nix, append "#foo" to the flake-uri.' >&2
    exit 1
  fi

  maybeRun mkdir -p "${mountPoint}"
  maybeRun chmod 755 "${mountPoint}" # bcachefs wants 755
  # shellcheck disable=SC2064
  if [[ -z ${dry_run-} ]]; then
     trap cleanupMountPoint EXIT
  fi

  outputs=$(nixBuild "${libexec_dir}"/install-cli.nix \
    "${nix_args[@]}" \
    --no-out-link \
    --impure \
    --argstr flake "$flake" \
    --argstr flakeAttr "$flakeAttr" \
    --argstr rootMountPoint "$mountPoint" \
    --arg writeEfiBootEntries "$writeEfiBootEntries" \
    --arg diskMappings "$(serialiaseArrayToNix diskMappings)" \
    --argstr extraSystemConfig "$extraSystemConfig" \
    -A installToplevel \
    -A closureInfo \
    -A "$diskoAttr")

  IFS=$'\n' mapfile -t artifacts <<<"$outputs"
  nixos_system=${artifacts[0]}
  closure_info=${artifacts[1]}
  disko_script=${artifacts[2]}

  if [[ -n ${dry_run-} ]]; then
    echo "Would run: $disko_script"
    echo "Would run: nixos-install --system '$nixos_system' --root '$mountPoint'"
    exit 0
  fi

  # We don't want swap as can break your running system in weird ways if you eject the disk
  # Hopefully disko-install has enough RAM to run without swap, otherwise we can make this configurable in future.
  DISKO_SKIP_SWAP=1 "$disko_script"

  for source in "${!extraFiles[@]}"; do
    destination=${extraFiles[$source]}
    mkdir -p "$mountPoint/$(dirname "$destination")"
    cp -ar "$source" "$mountPoint/$destination"
  done

  # nix copy uses up a lot of memory and we work around issues with incorrect checksums in our store
  # that can be caused by using closureInfo in combination with multiple builders and non-deterministic builds.
  # Therefore if we have a blank store, we copy the store paths and registration from the closureInfo.
  if [[ ! -f "${mountPoint}/nix/var/nix/db/db.sqlite" ]]; then
    echo "Copying store paths" >&2
    mkdir -p "${mountPoint}/nix/store"
    xargs cp --recursive --target "${mountPoint}/nix/store"  < "${closure_info}/store-paths"
    echo "Loading nix database" >&2
    NIX_STATE_DIR=${mountPoint}/nix/var/nix nix-store --load-db < "${closure_info}/registration"
  fi

  nixos-install --no-channel-copy --no-root-password --system "$nixos_system" --root "$mountPoint"
}

if main "$@"; then
  echo "disko-install succeeded"
  exit 0
else
  echo "disko-install failed" >&2
  exit 1
fi