#!/usr/bin/env bash

# defaults
FLAKE_URI="."
CONFIG_FILE="./deploy.json"
ACTION="switch"
USE_SUDO=true
DO_BUILD=true
FILTER_HOSTS=()

usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS] [ACTION]

Arguments:
  ACTION                 switch | boot | test (Default: $ACTION)

Options:
  -n, --host NAME        Deploy only this host (can be repeated)
  -f, --flake URI        URI of the flake (Default: $FLAKE_URI)
  -c, --config FILE      Deployment config file (Default: $CONFIG_FILE)
  --no-sudo              Do not pass sudo-related flags to nixos-rebuild.
  --skip-build           Skip the explicit 'build' step before deployment.
  -h, --help             Show this help.

Examples:
  $(basename "$0")                      Deploy all hosts
  $(basename "$0") -n cryodev-main      Deploy only cryodev-main
  $(basename "$0") -n host-a -n host-b  Deploy host-a and host-b
EOF
}

_status() { echo -e "\033[0;34m> $1\033[0m"; }
success() { echo -e "\033[0;32m$1\033[0m"; }
error()   { echo -e "\033[0;31mError: $1\033[0m" >&2; exit 1; }

while [[ $# -gt 0 ]]; do
  case "$1" in
    switch|boot|test)  ACTION="$1"; shift ;;
    -n|--host)         FILTER_HOSTS+=("$2"); shift 2 ;;
    -f|--flake)        FLAKE_URI="$2"; shift 2 ;;
    -c|--config)       CONFIG_FILE="$2"; shift 2 ;;
    --no-sudo)         USE_SUDO=false; shift ;;
    --skip-build)      DO_BUILD=false; shift ;;
    -h|--help)         usage; exit 0 ;;
    *)                 error "Invalid argument '$1'" ;;
  esac
done

command -v jq &> /dev/null || error "jq is not installed."
[ -f "$CONFIG_FILE" ] || error "Config '$CONFIG_FILE' not found."

BUILD_HOST=$(jq -r '.buildHost // "localhost"' "$CONFIG_FILE")
[[ "$BUILD_HOST" =~ ^(127\.0\.0\.1|::1)$ ]] && BUILD_HOST="localhost"

SSH_PORT=$(jq -r '.sshPort // "22"' "$CONFIG_FILE")
export NIX_SSHOPTS="-p $SSH_PORT"

mapfile -t ALL_ENTRIES < <(jq -r '.hosts[] | "\(.name) \(.address)"' "$CONFIG_FILE")
[ ${#ALL_ENTRIES[@]} -eq 0 ] && error "No hosts defined in $CONFIG_FILE"

# Filter hosts if -n was provided
HOST_ENTRIES=()
if [ ${#FILTER_HOSTS[@]} -gt 0 ]; then
  for entry in "${ALL_ENTRIES[@]}"; do
    read -r name _address <<< "$entry"
    for filter in "${FILTER_HOSTS[@]}"; do
      if [[ "$name" == "$filter" ]]; then
        HOST_ENTRIES+=("$entry")
        break
      fi
    done
  done
  # Check for unknown hosts
  for filter in "${FILTER_HOSTS[@]}"; do
    found=false
    for entry in "${ALL_ENTRIES[@]}"; do
      read -r name _ <<< "$entry"
      [[ "$name" == "$filter" ]] && found=true && break
    done
    [[ "$found" == false ]] && error "Host '$filter' not found in $CONFIG_FILE"
  done
  [ ${#HOST_ENTRIES[@]} -eq 0 ] && error "No matching hosts found"
else
  HOST_ENTRIES=("${ALL_ENTRIES[@]}")
fi

echo "Action:      $ACTION"
echo "Flake:       $FLAKE_URI"
echo "Builder:     $BUILD_HOST"
echo "SSH Port:    $SSH_PORT"
echo "Hosts:       $(printf '%s ' "${HOST_ENTRIES[@]}" | sed 's/ [^ ]*//g; s/  */, /g')"

if [ "$DO_BUILD" = true ]; then
  _status "Building configurations..."
  for entry in "${HOST_ENTRIES[@]}"; do
    read -r name address <<< "$entry"
    echo "------------------------------------------------"
    echo "Building host '$name':"

    CMD=("nixos-rebuild" "build" "--flake" "${FLAKE_URI}#${name}")
    [[ "$BUILD_HOST" != "localhost" ]] && CMD+=("--build-host" "$BUILD_HOST")

    "${CMD[@]}" || error "Build failed for $name"
    success "Build for host '$name' successful."
  done
fi

_status "Deploying to targets..."
for entry in "${HOST_ENTRIES[@]}"; do
  read -r name address <<< "$entry"
  echo "------------------------------------------------"
  echo "Deploying to host '$name' ($address):"

  CMD=("nixos-rebuild" "$ACTION" "--flake" "${FLAKE_URI}#${name}" "--target-host" "$address")
  [[ "$BUILD_HOST" != "localhost" ]] && CMD+=("--build-host" "$BUILD_HOST")
  [[ "$USE_SUDO" = true ]] && CMD+=("--sudo" "--ask-sudo-password")

  "${CMD[@]}" || error "Activation failed for $name"
  success "Host '$name' updated."
done

success "Deployment complete."