# NixOS and standalone Home Manager rebuild script

# Defaults
FLAKE_PATH="${REBUILD_FLAKE_PATH:-${FLAKE_PATH:-$HOME/.config/nixos}}" # Default flake path
USER="${REBUILD_USER:-$(whoami)}"                                      # Default username
HOST="${REBUILD_HOST:-$(hostname)}"                                    # Default hostname
BUILD_HOST="${REBUILD_BUILD_HOST:-}"                                   # Default build host
TARGET_HOST="${REBUILD_TARGET_HOST:-}"                                 # Default target host
UPDATE="${REBUILD_UPDATE:-0}"                                          # Default to not update
UPDATE_INPUTS="${REBUILD_UPDATE_INPUTS:-}"                             # Default list of inputs
ROLLBACK="${REBUILD_ROLLBACK:-0}"                                      # Default to not rollback
SHOW_TRACE="${REBUILD_SHOW_TRACE:-0}"                                  # Default to not show trace

# Output functions
_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; }

# Function to display the help message
Help() {
  echo "Wrapper script for 'nixos-rebuild switch' and 'home-manager switch' commands."
  echo "Usage: rebuild <command> [OPTIONS]"
  echo
  echo "Commands:"
  echo "  nixos                  Rebuild NixOS configuration"
  echo "  home                   Rebuild Home Manager configuration"
  echo "  all                    Rebuild both NixOS and Home Manager configurations"
  echo "  help                   Show this help message"
  echo
  echo "Options (for NixOS and Home Manager):"
  echo "  -H, --host <host>      Specify the hostname (as in 'nixosConfiguraions.<host>'). Default: $HOST"
  echo "  -p, --path <path>      Set the path to the flake directory. Default: $FLAKE_PATH"
  echo "  -U, --update [inputs]  Update all flake inputs. Optionally provide comma-separated list of inputs to update instead."
  echo "  -r, --rollback         Don't build the new configuration, but use the previous generation instead"
  echo "  -t, --show-trace       Show detailed error messages"
  echo
  echo "NixOS only options:"
  echo "  -B, --build-host <user@example.com>   Use a remote host for building the configuration via SSH"
  echo "  -T, --target-host <user@example.com>  Deploy the configuration to a remote host via SSH. If '--host' is specified, it will be used as the target host."
  echo
  echo "Home Manager only options:"
  echo "  -u, --user <user>      Specify the username (as in 'homeConfigurations.<user>@<host>'). Default: $USER"
}

# Function to rebuild NixOS configuration
Rebuild_nixos() {
  local FLAKE="$FLAKE_PATH#$HOST"

  # Construct rebuild command
  local CMD=("nixos-rebuild" "switch" "--sudo")
  CMD+=("--flake" "$FLAKE")
  [ "$ROLLBACK" = 1 ] && CMD+=("--rollback")
  [ "$SHOW_TRACE" = 1 ] && CMD+=("--show-trace")
  [ -n "$BUILD_HOST" ] && CMD+=("--build-host" "$BUILD_HOST")
  if [ "$HOST" != "$(hostname)" ] && [ -z "$TARGET_HOST" ]; then
    TARGET_HOST="$HOST"
    _status "Using '$TARGET_HOST' as target host."
  fi
  [ -n "$TARGET_HOST" ] && CMD+=("--target-host" "$TARGET_HOST")
  [[ -n "$TARGET_HOST" || -n "$BUILD_HOST" ]] && CMD+=("--ask-sudo-password")

  # Build config first so we can diff it
  local BUILD_CMD=("nixos-rebuild" "build" "--no-build-output" "--flake" "$FLAKE")
  [ "$SHOW_TRACE" = 1 ] && BUILD_CMD+=("--show-trace")
  [ -n "$BUILD_HOST" ] && BUILD_CMD+=("--build-host" "$BUILD_HOST")

  if [ "$ROLLBACK" = 0 ]; then
    _status "Building NixOS configuration '$FLAKE'..."
    _status "Executing command: ${BUILD_CMD[*]}"
    "${BUILD_CMD[@]}" || error "NixOS build failed"

    _status "Switching to new NixOS configuration"
  else
    _status "Rolling back to last NixOS generation"
  fi

  sudo -v
  _status "Executing command: ${CMD[*]}"
  "${CMD[@]}" || error "NixOS rebuild failed"
  success "NixOS rebuild completed successfully."
}

# Function to rebuild Home Manager configuration
Rebuild_home() {
  local FLAKE="$FLAKE_PATH#$USER@$HOST"

  if [ -n "$BUILD_HOST" ] || [ -n "$TARGET_HOST" ]; then
    error "Remote building is not supported for Home Manager."
  fi

  # Construct rebuild command
  local CMD=()
  if [ "$ROLLBACK" = 1 ]; then
    local rollback_path
    rollback_path=$(home-manager generations | sed -n '2p' | grep -o '/nix/store[^ ]*')
    CMD+=("$rollback_path/activate")
  else
    CMD=("home-manager" "switch" "--flake" "$FLAKE")
    [ "$SHOW_TRACE" = 1 ] && CMD+=("--show-trace")
  fi

  # Build config first so we can diff it
  if [ "$ROLLBACK" = 0 ]; then
    local BUILD_CMD=("home-manager" "build" "--no-out-link" "--flake" "$FLAKE")
    [ "$SHOW_TRACE" = 1 ] && BUILD_CMD+=("--show-trace")
    _status "Building Home Manager configuration '$FLAKE'..."
    _status "Executing command: ${BUILD_CMD[*]}"
    "${BUILD_CMD[@]}" || error "Home Manager build failed"

    _status "Switching to new Home Manager configuration"
  else
    _status "Rolling back to last Home Manager generation"
  fi

  _status "Executing command: ${CMD[*]}"
  "${CMD[@]}" || error "Home Manager rebuild failed"
  success "Home Manager rebuild completed successfully."
}

# Function to update flake repositories
Update() {
  _status "Updating flake inputs..."

  # Construct update command as an array
  local CMD=("nix" "flake" "update" "--flake" "$FLAKE_PATH")
  if [ -n "$UPDATE_INPUTS" ]; then
    # Split comma-separated inputs and pass them to nix flake update
    IFS=',' read -ra INPUTS <<< "$UPDATE_INPUTS"
    for input in "${INPUTS[@]}"; do
      CMD+=("$input")
    done
  fi

  _status "Executing command: ${CMD[*]}"
  "${CMD[@]}" || error "Failed to update flake repositories"
  success "Flake repositories updated successfully."
}

# Parse command-line options
if [[ -z "${1:-}" ]]; then
  echo -e "\033[0;31mError: No command specified. Printing help page.\033[0m" >&2
  Help
  exit 1
fi

COMMAND=$1
shift

# Handle help command early
if [ "$COMMAND" = "help" ] || [ "$COMMAND" = "--help" ] || [ "$COMMAND" = "-h" ]; then
  Help
  exit 0
fi

while [ $# -gt 0 ]; do
  case "${1:-}" in
    -H|--host)
      if [ -n "${2:-}" ]; then
        HOST="$2"
        shift 2
      else
        error "-H|--host option requires an argument"
      fi
      ;;
    -u|--user)
      if [ -n "${2:-}" ]; then
        USER="$2"
        shift 2
      else
        error "-u|--user option requires an argument"
      fi
      ;;
    -p|--path)
      if [ -n "${2:-}" ]; then
        FLAKE_PATH="$2"
        shift 2
      else
        error "-p|--path option requires an argument"
      fi
      ;;
    -U|--update)
      UPDATE=1
      # Check if next argument is a non-option
      if [ $# -gt 1 ] && [ "${2#-}" = "${2:-}" ]; then
        UPDATE_INPUTS="$2"
        shift 2
      else
        shift
      fi
      ;;
    -r|--rollback)
      ROLLBACK=1
      shift
      ;;
    -t|--show-trace)
      SHOW_TRACE=1
      shift
      ;;
    -B|--build-host)
      if [ -n "${2:-}" ]; then
        BUILD_HOST="$2"
        shift 2
      else
        error "-B|--build-host option requires an argument"
      fi
      ;;
    -T|--target-host)
      if [ -n "${2:-}" ]; then
        TARGET_HOST="$2"
        shift 2
      else
        error "-T|--target-host option requires an argument"
      fi
      ;;
    *)
      error "Unknown option '$1'"
      ;;
  esac
done

# Check if script is run with sudo
if [ "$EUID" -eq 0 ]; then
  error "Do not run this script with sudo."
fi

# Check if flake path exists
if [ ! -d "$FLAKE_PATH" ]; then
  error "Flake path '$FLAKE_PATH' does not exist"
fi

# Ignore trailing slash in flake path
FLAKE_PATH="${FLAKE_PATH%/}"

# Check if flake.nix exists
if [ ! -f "$FLAKE_PATH/flake.nix" ]; then
  error "flake.nix does not exist in '$FLAKE_PATH'"
fi

# Execute updates and rebuilds based on the command
[ "$UPDATE" = 1 ] && Update

case "$COMMAND" in
  nixos|os)
    Rebuild_nixos
    ;;
  home|hm)
    Rebuild_home
    ;;
  all)
    Rebuild_nixos
    Rebuild_home
    ;;
  *)
    error "Unknown command '$COMMAND'"
    ;;
esac