#!/usr/bin/env bash

unset GREP_COLOR
unset GREP_OPTIONS

source "$rvm_scripts_path/base"

__error_on_result() {
  if [[ "$1" -gt 0 ]]; then
    "$rvm_scripts_path/log" "fail" "$2 - Aborting now."
    return 0
  else
    return 1
  fi
}

snapshot_save() {

  if [[ -z "$1" ]]; then
    echo "Usage: rvm snapshot save name" >&2
    echo "Saves a snapshot to <name>.tar.gz in the current directory." >&2
    return 1
  fi

  # Create the temporary directory.
  local snapshot_temp_path="$rvm_tmp_path/$$-snapshot"
  \rm -rf "$snapshot_temp_path"
  \mkdir -p "$snapshot_temp_path"

  "$rvm_scripts_path/log" "info" "Backing up a list of aliases"
  \cp "$rvm_config_path/alias" "$snapshot_temp_path/"

  "$rvm_scripts_path/log" "info" "Backing up your user preferences"
  \cp "$rvm_config_path/user" "$snapshot_temp_path/"

  "$rvm_scripts_path/log" "info" "Backing up your installed packages"
  sed -e 's/-//' -e 's/^lib//' < "$rvm_config_path/packages" | awk -F= '{print $1}' | sort | uniq > "$snapshot_temp_path/packages"

  "$rvm_scripts_path/log" "info" "Backing up all of your gemsets"
  \mkdir -p "$snapshot_temp_path/gems"
  __rvm_pushpop "$snapshot_temp_path/gems"
  for snapshot_gemset in $("$rvm_scripts_path/list" gemsets strings) ; do
    __rvm_become "$snapshot_gemset"
    result="$?"
    __error_on_result "$result" "Error becoming ruby $snapshot_gemset" && return "$result"
    "$rvm_scripts_path/gemsets" export "${snapshot_gemset}.gems" >/dev/null
    result="$?"
    __error_on_result "$result" "Error exporting gemset contents for $snapshot_gemset" && return "$result"
    \mkdir -p "./$snapshot_gemset/"
    [[ -d "$GEM_HOME/cache/" ]] && \cp -R "$GEM_HOME/cache/" "./$snapshot_gemset/"
  done; unset snapshot_gemset
  __rvm_pushpop

  "$rvm_scripts_path/log" "info" "Backing up all of your installed rubies"
  echo '#!/usr/bin/env bash -e' > "$snapshot_temp_path/install-rubies.sh"
  echo "source \"\$rvm_scripts_path/rvm\" || true" >> "$snapshot_temp_path/install-rubies.sh"
  local snapshot_ruby_name_file="$rvm_tmp_path/$$-rubies"
  local snapshot_alias_name_file="$rvm_tmp_path/$$-aliases"
  local snapshot_installable_file="$rvm_tmp_path/$$-installable"

  "$rvm_scripts_path/alias" list | awk -F ' => ' '{print $1}' | sort | uniq 2>/dev/null > "$snapshot_alias_name_file"
  "$rvm_scripts_path/list" strings | \tr ' ' '\n' | sort | uniq > "$snapshot_ruby_name_file"
  comm -2 -3 "$snapshot_ruby_name_file" "$snapshot_alias_name_file" > "$snapshot_installable_file"
  \rm -rf "$snapshot_ruby_name_file" "$snapshot_alias_name_file"

  local snapshot_primary_ruby="$(\grep '^\(ree\|ruby-1.8.7\)' < "$snapshot_installable_file" | \grep -v '-head$' | sort -r | head -n1)"
  local snapshot_ruby_order="$snapshot_primary_ruby $(\grep -v "$snapshot_primary_ruby" < "$snapshot_installable_file")"
  for snapshot_ruby_name in $snapshot_ruby_order ; do
    snapshot_install_command="$(__rvm_recorded_install_command "$snapshot_ruby_name")"
    if [[ -n "$snapshot_install_command" ]]; then
      echo "rvm install $snapshot_install_command" | sed "s#$rvm_path#'\\\"\$rvm_path\\\"'#" >> "$snapshot_temp_path/install-rubies.sh"
    else
      __rvm_become "$snapshot_ruby_name"
      ruby "$rvm_path/lib/rvm/install_command_dumper.rb" >> "$snapshot_temp_path/install-rubies.sh"
    fi
    unset snapshot_install_command
  done; unset snapshot_ruby_name snapshot_primary_ruby

  \rm -rf "$snapshot_installable_file"

  "$rvm_scripts_path/log" "info" "Compressing snapshotting"
  local destination_path="$PWD"
  __rvm_pushpop "$snapshot_temp_path"
  \rm -rf "$destination_path/$1.tar.gz"
  tar czf "$destination_path/$1.tar.gz" .
  result="$?"
  __error_on_result "$result" "Error creating archive $destination_path/$1.tar.gz" && return "$result"
  __rvm_pushpop

  "$rvm_scripts_path/log" "info" "Cleaning up"
  \rm -rf "$snapshot_temp_path"

  "$rvm_scripts_path/log" "info" "Snapshot complete"
}

snapshot_load() {
  if [[ -z "$1" ]]; then
    echo "Usage: rvm snapshot load name" >&2
    echo "Loads a snapshot from <name>.tar.gz in the current directory." >&2
    return 1
  fi

  local snapshot_archive="$PWD/$(echo "$1" | sed 's/.tar.gz$//').tar.gz"

  if ! [[ -s "$snapshot_archive" ]]; then
    echo "The provides snapshot '$(basename "$snapshot_archive")' doesn't exist." >&2
    return 1
  fi

  local snapshot_temp_path="$rvm_tmp_path/$$-snapshot"

  \rm -rf "$snapshot_temp_path"
  \mkdir -p "$snapshot_temp_path"

  "$rvm_scripts_path/log" "info" "Extracting snapshot"
  __rvm_pushpop "$snapshot_temp_path"
  tar xzf "$snapshot_archive"
  result="$?"
  __error_on_result "$result" "Error extracting the archive '$snapshot_archive'" && return "$result"
  __rvm_pushpop

  "$rvm_scripts_path/log" "info" "Restoring user settings"
  \cp -f "$snapshot_temp_path/user" "$rvm_config_path/user"

  "$rvm_scripts_path/log" "info" "Installing rvm-managed packages"
  for snapshot_package in $(cat "$snapshot_temp_path/packages"); do
    "$rvm_scripts_path/package" install "$snapshot_package"
    result="$?"
    __error_on_result "$result" "Error installing package '$snapshot_package'" && return "$result"
  done; unset snapshot_package

  "$rvm_scripts_path/log" "info" "Installing rubies"
  chmod +x "$snapshot_temp_path/install-rubies.sh"
  "$snapshot_temp_path/install-rubies.sh"
  result="$?"
  __error_on_result "$result" "Error importing rubies." && return "$result"

  export rvm_create_flag=1
  "$rvm_scripts_path/log" "info" "Setting up gemsets"
  __rvm_pushpop "$snapshot_temp_path/gems"
  for snapshot_gemset in $(\ls | \grep '\.gems$' | sed 's/.gems$//'); do
    __rvm_become "$snapshot_gemset"
    result="$?"
    __error_on_result "$result" "Error becoming '$snapshot_gemset'" && return "$result"
    \mkdir -p "$GEM_HOME/cache/"
    \cp -Rf "$snapshot_gemset/" "$GEM_HOME/cache/"
    result="$?"
    __error_on_result "$result" "Error copying across cache for $snapshot_gemset" && return "$result"
    "$rvm_scripts_path/gemsets" import "$snapshot_gemset" >/dev/null 2>&1
    result="$?"
    __error_on_result "$result" "Error importing gemset for $snapshot_gemset" && return "$result"
  done; unset snapshot_gemset
  __rvm_pushpop

  "$rvm_scripts_path/log" "info" "Restoring aliases"
  while read -r package_info; do
    # Note: this assumes an '=' int the input...
    local alias_name="${package_info/=*}"
    local alias_ruby="${package_info/*=}"

    "$rvm_scripts_path/alias" create "$alias_name" "$alias_ruby"
    if [[ "$alias_name" = "default" ]]; then
      (source "$rvm_scripts_path/rvm" && rvm use "$alias_ruby" --default) >/dev/null 2>&1
      result="$?"
      __error_on_result "$result" "Error setting default to $alias_ruby" && return "$result"
    fi
  done < "$snapshot_temp_path/alias"
  unset package_info

  "$rvm_scripts_path/log" "info" "Cleaning up load process"
  \rm -rf "$snapshot_temp_path"

  "$rvm_scripts_path/log" "info" "Loaded snapshot from $(basename "$snapshot_archive")"
}

snapshot_usage() {
  echo "Usage: rvm snapshot {save,load} file" >&2
  return 1
}

action="$1"
[[ "$#" -gt 0 ]] && shift

case "$action" in
  save) snapshot_save "$@" ;;
  load) snapshot_load "$@" ;;
  *)    snapshot_usage ;;
esac

exit $?