#!/usr/bin/env zsh

# I don't like the user story for Nix's binaries, unlike Guix's.  So I wrote
# this to abstract away the unneeded complexity for everyday maintenance of my
# NixOS config.
#
#   hey cd|check
#     nix flake check $DOTFILES
#   hey sh|show
#     nix flake show $DOTFILES
#   hey u|update [INPUT...]
#     nix flake update --recreate-lock-file $DOTFILES
#   hey re|rebuild [ARGS...]
#     nixos-rebuild --flake $DOTFILES#$HOST ${@:-build}
#   hey t|test        -> nixos-rebuild --flake $DOTFILES#$HOST --fast test
#   hey theme $THEME  -> THEME=$THEME hey test (quick theme swapping)
#   hey rollback      -> hey rebuild --rollback switch
#   hey s|search      -> nix search nixpkgs $@
#   hey swap $PATH    -> swaps all specified nix-store symlinks (or recursively
#                        in a specified directory) with copies so they can be
#                        modified in place (while tweaking). Run again to
#                        restore the old link (careful, you'll lose work if you
#                        don't save your changes elsewhere!)
#   hey vm            -> hey rebuild vm
#   hey up|upgrade    -> hey update && hey rebuild switch
#   hey gc            -> GC user profile, optimise nix store
#   sudo hey gc       -> GC, optimise nix store, & delete old boot entries
#   hey push REMOTE   -> 1. Build machine locally, 2. Copy result to remote, 3. Activate config on remote
#   hey gen|generations -> explore, manage, diff across generations (WIP)
#
# hey (with no sub-command) can also be used as an alternative to nix-env:
#
#   hey -q
#   hey -iA nixos.htop
#   hey -e htop

set -e
if [[ -z $FLAKE ]]; then
  for dir in /etc/nixos /etc/dotfiles; do
    if [[ -f "$dir/flake.nix" ]]; then
      FLAKE=$dir
      break
    fi
  done
fi

#
function main {
  function _do { [[ $DRYRUN ]] && echo "$@" || "$@" }

  ### Helpers
  function _usage {
    echo "Usage: hey COMMAND [FLAGS...] [ARGS...]"
    echo
    echo "Available commands:"
    echo "  check (ch)         Same as 'nix flake check \$DOTFILES'"
    echo "  gc                 Garbage collect and optimize store"
    echo "  generations (gen)  TODO"
    echo "  push               Push thie config to a remote NixOS system and rebuild it"
    echo "  rebuild (re)       Rebuild current system flake"
    echo "  rollback           Roll back to last generation"
    echo "  search (s)         Searches nixpkgs for a package"
    echo "  show (sh)          Same as 'nix flake show \$DOTFILES'"
    echo "  swap               Swap nix-store link (or directory thereof) for tweaking"
    echo "  test (t)           Quickly rebuild, for quick iteration"
    echo "  theme              Quickly swap to another theme"
    echo "  upgrade (up)       Update flake lockfile and rebuild"
    echo "  vm                 Build flake into a VM image"
  }

  function _rebuild {
    # Move to $DOTFILES to prevent polluting the system with result links
    pushd $FLAKE >/dev/null
    # REVIEW Use --impure instead when it's in a stable release of nix
    _do ${SUDO:+sudo} nixos-rebuild \
      --flake "${FLAKE}#${HOSTNAME:-$HOST}" \
      --option pure-eval no \
      ${@:-build}
    popd >/dev/null
  }

  function _update {
    echo "Updating NixOS flake"
    if [[ "$@" ]]; then
      for flake in $@; do
        _do nix flake update --update-input "$flake" "$FLAKE"
      done
    else
      _do nix flake update --recreate-lock-file "$FLAKE"
    fi
  }

  function _gc {
    if [[ $EUID == 0 ]]; then
      echo "Cleaning up your system..."
      _do nix-collect-garbage ${@:--d}
      _do nix-store --optimise
      # Clear out left-over secrets. The switch-to-configuration call below will
      # regenerate them.
      [[ -d /run/secrets ]] && _do rm -f /run/secrets/*
      # nix-collect-garbage is a Nix tool, not a NixOS tool. It won't delete old
      # boot entries until you do a nixos-rebuild (which means we'll always have
      # 2 boot entries at any time). Instead, we properly delete them by
      # reloading the current environment.
      _do nix-env --delete-generations old --profile /nix/var/nix/profiles/system
      _do /nix/var/nix/profiles/system/bin/switch-to-configuration switch
    else
      echo "Cleaning up your user profile..."
      _do nix-collect-garbage ${@:--d}
    fi
  }

  function _swap {
    for target in $@; do
      if [ -d "$target" ]; then
        files=( "$target"/**/*.nix-store-backup )
        links=( $(find "$target" -type l) )
        if [[ "$files" ]]; then
          echo "There are swapped files here. Unswapping!"
          for file in $files[@]; do
            _swap "$link"
          done
        else
          echo "Swapping files in $target"
          for link in $links[@]; do
            if [[ "$(readlink $link)" == /nix/store/* ]]; then
              _swap "$link"
            fi
          done
        fi
      elif [ -f "$target.nix-store-backup" ]; then
        echo "Unswapping $target"
        _do mv -i "$target.nix-store-backup" "$target"
      elif [ -f "$target" ]; then
        echo "Swapping $target"
        _do mv "$target" "$target.nix-store-backup"
        _do cp "$target.nix-store-backup" "$target"
      else
        >&2 echo "Not a valid target: $target"
        exit 1
      fi
    done
  }

  ###
  eval set -- $(getopt -n hey -o df:h: -l dryrun,flake:,host: -- "$@")
  while :
  do
    case "$1" in
      -d|--dryrun) DRYRUN=1;    shift   ;;
      -f|--flake)  FLAKE=$2;    shift 2 ;;
      -h|--host)   HOSTNAME=$2; shift 2 ;;
      --) shift; break ;;
    esac
  done
  if [[ ! -d $FLAKE ]]; then
    >&2 echo "Couldn't find flake.nix in $FLAKE"
    exit 1
  fi

  cmd="$1"
  [[ $# > 0 ]] && shift
  case "$cmd" in
    sh|shell)    _do nix develop "$FLAKE" "$@" ;;
    repl)        _do nix repl "$FLAKE/flake.nix" "<nixpkgs>" "$@" ;;

    ch|check)    _do nix flake check "$@" "$FLAKE"      ;;
    sh|show)     _do nix flake show "$@" "$FLAKE"       ;;

    theme)       THEME="$1" SUDO=1 _rebuild --fast test ;;
    t|test)      SUDO=1 _rebuild --fast test "$@"       ;;
    re|rebuild)  SUDO=1 _rebuild "${@:-switch}"         ;;
    rollback)    SUDO=1 _rebuild --rollback switch      ;;
    vm)          SUDO=1 _rebuild "$@" vm                ;;

    u|update)    _update "$@"                           ;;
    up|upgrade)  _update && SUDO=1 _rebuild switch      ;;
    s|search)    nix search nixpkgs "$@"                ;;

    swap)        _swap "$@" ;;
    gc)          _gc "$@" ;;
    push)
      target=$1
      cmd=${2:-test}
      echo "Enabling root ssh access on $target"
      if _do ssh -t "$target" "sudo mkdir -p /root/.ssh && sudo cp -f ~/.ssh/authorized_keys /root/.ssh/authorized_keys"; then
        trap '_do ssh -t "root@$target" "rm -f /root/.ssh/authorized_keys"' EXIT
        echo "Building $target ($cmd)"
        HOSTNAME="${HOSTNAME:-$target}" _rebuild \
          --build-host localhost \
          --target-host "root@$target" \
          $cmd
        echo "Undoing root ssh access on $target"
      fi
      ;;

    gen|generations)
      profile=/nix/var/nix/profiles/system
      case "$1" in
        switch) >&2 echo "Targeted rollback not supported yet" ;;  # TODO Is this true?
        diff)
          to=/nix/var/nix/profiles/system-${2}-link
          diff \
            <(nix-store -q --references $profile) \
            <(nix-store -q --references $to)
          ;;
        remove)
          shift
          _do sudo nix-env --delete-generations --profile $profile $@
          ;;
        "")
          _do sudo nix-env --list-generations --profile $profile
          ;;
      esac
      ;;

    h|help|-h|--help) _usage ;;
    -*) _do nix-env "$cmd" "$@" ;;
    "") _usage ;;
    *)  echo "Couldn't find command: $cmd" 1>&2
        exit 1
        ;;
  esac
}

main "$@"