#!/usr/bin/env bash
#
# Description: Krathalan's Borg Helper: makes managing borg backups and repos
#              braindead easy with JSON configuration.

# -----------------------------------------
# -------------- Guidelines ---------------
# -----------------------------------------

# This script follows the Google Shell Style Guide:
# https://google.github.io/styleguide/shell.xml

# This script uses shellcheck: https://www.shellcheck.net/

# See https://vaneyckt.io/posts/safer_bash_scripts_with_set_euxo_pipefail/
set -Eeuo pipefail

# -----------------------------------------
# ----------- Program variables -----------
# -----------------------------------------

# Colors
GREEN=$(tput bold && tput setaf 2)
RED=$(tput bold && tput setaf 1)
WHITE=$(tput sgr0 && tput bold)
YELLOW=$(tput sgr0 && tput setaf 3)
GREY=$(tput sgr0 && tput setaf 8)
NC=$(tput sgr0) # No color/turn off all tput attributes

exit_script_on_failure()
{
  printf "%sError%s: %s\n" "${RED}" "${NC}" "$1" >&2

  exit 1
}

# Script (self) information
SCRIPT_NAME="${0##*/}"

# Check for config file
readonly CONFIG_FILE_DIRECTORY="${XDG_CONFIG_HOME:-${HOME}/.config}/krathalan"
readonly CONFIG_FILE="${CONFIG_FILE_DIRECTORY}/${SCRIPT_NAME}.json"

if [[ ! -f "${CONFIG_FILE}" ]]; then
  exit_script_on_failure "No configuration file found at ${CONFIG_FILE}"
fi

# Load settings
KBH_SETTINGS_JSON="$(<"${CONFIG_FILE}")"
KBH_REPOS_NUMBER="$(jaq -r ".repos | length" <<< "${KBH_SETTINGS_JSON}")"
KBH_ARCHIVE_NAME="$(jaq -r .borg_archive_name_command <<< "${KBH_SETTINGS_JSON}")"

if [[ "${KBH_ARCHIVE_NAME}" == "null" ]]; then
  # Use default archive name
  KBH_ARCHIVE_NAME="date +%Y-%m-%d-%H%M%S"
fi

readonly KBH_SETTINGS_JSON KBH_REPOS_NUMBER KBH_ARCHIVE_NAME SCRIPT_NAME GREEN RED WHITE NC

# -----------------------------------------
# --------------- Functions ---------------
# -----------------------------------------

print_help()
{
  printf "%s" "\
${SCRIPT_NAME} - Krath's Borg Helper

More detailed usage information, including configuration and override
variables, is found in the man page.

=> help|-h|--help - show this help
=> backup         - backs up all backup directories
=> delete         - deletes archive from repo
=> info           - lists borg repos general information
=> list           - lists archives from borg repos
=> init           - initializes borg repos
=> mount          - mounts repo or archive from specified repo onto /mnt/borg
=> prune          - prunes archives according to config setting
=> status         - displays config and status of all repos
"

  exit 0
}

check_ssh_key()
{
  # Load settings
  local -r sshKeyPasswordCopyCommand="$(jaq -r .ssh_key_password_copy_command <<< "${KBH_SETTINGS_JSON}")"

  if [[ "${sshKeyPasswordCopyCommand}" == "null" ]]; then
    return
  fi

  if ! ssh-add -L &> /dev/null; then
    print_script_message "You need to add your ssh key. Copying passphrase and initiating adding..."
    
    ${sshKeyPasswordCopyCommand}

    ssh-add
  fi
}

exit_script_on_failure()
{
  printf "%sError%s: %s\n" "${RED}" "${NC}" "$1" >&2

  exit 1
}

print_script_message()
{
  printf "%s==>%s %s%s\n" "${GREEN}" "${WHITE}" "$1" "${NC}"
}

get_mount_point()
{
  local -r mountPointSetting="$(jaq -r .mount_point <<< "${KBH_SETTINGS_JSON}")"

  # Discern where to mount
  # If KBH_MOUNTPOINT env variable is specified, use that first.
  # Otherwise use value of .mount_point in conf file.
  # Otherwise default to /mnt/borg.
  if [[ -n "${KBH_MOUNTPOINT:-}" ]]; then
    local -r mountPointOfficial="${KBH_MOUNTPOINT}"
  elif [[ -n "${mountPointSetting}" ]]; then
    local -r mountPointOfficial="${mountPointSetting}"
  else
    local -r mountPointOfficial="/mnt/borg"
  fi

  printf "%s" "${mountPointOfficial}"
}

# $1: name (or alternative name) of repo to return json data of
select_repo_by_name()
{
  # Attempt to match repo from name
  local repoData
  repoData="$(jaq -r ".repos[] | select(.name==\"$1\")" <<< "${KBH_SETTINGS_JSON}")"

  # If not found, attempt to match repo from alternative names
  if [[ -z "${repoData:-}" ]]; then
    repoData="$(jaq -r ".repos[] | select(.alternative_names[]==\"$1\")" <<< "${KBH_SETTINGS_JSON}")"
  fi

  # Crash if repo not found
  if [[ -z "${repoData:-}" ]]; then
    exit_script_on_failure "Could not locate repo with name or alternative name \"$1\""
  fi

  printf "%s" "${repoData}"
}

# $1: index of repo to return json data of
select_repo_by_index()
{
  local repoData
  repoData="$(jaq -r ".repos[$1]" <<< "${KBH_SETTINGS_JSON}")"

  if [[ "${repoData}" == "null" ]]; then
    exit_script_on_failure "Could not locate repo with index $1 in config file"
  fi

  printf "%s" "${repoData}"
}

repo_info_helper()
{
  local -r repoName="$(jaq -r .name <<< "$1")"
  local -r repoPath="$(jaq -r .path <<< "$1")"
  local -r passwordCommand="$(jaq -r .password_command <<< "$1")"
  
  print_script_message "Repo info on ${repoName}..."

  BORG_PASSPHRASE="$(${passwordCommand})" borg info "${repoPath}"
}

# $*: repos (optional)
repo_info()
{
  shift
  local repoData=""

  # If repo specified, print info for only that repo
  if [[ $# -gt 0 ]]; then
    while [[ $# -gt 0 ]]; do
      repoData="$(select_repo_by_name "$1")"
      repo_info_helper "${repoData}"

      printf "\n"
      shift
    done
  else
    # If repo not specified, print info for all repos
    print_script_message "Listing repo info."

    local whileCounter=0
    while [[ "${whileCounter}" -lt "${KBH_REPOS_NUMBER}" ]]; do
      repoData="$(select_repo_by_index "${whileCounter}")"
      repo_info_helper "${repoData}"

      whileCounter="$(( whileCounter + 1 ))"
      printf "\n"
    done
  fi
}

list_repo_archives_helper()
{
  local -r repoName="$(jaq -r .name <<< "$1")"
  local -r repoPath="$(jaq -r .path <<< "$1")"
  local -r passwordCommand="$(jaq -r .password_command <<< "$1")"
  
  print_script_message "Repo archives on ${repoName}..."

  BORG_PASSPHRASE="$(${passwordCommand})" borg list "${repoPath}"
}

# $*: repos (optional)
list_repo_archives()
{
  shift
  local repoData=""

  # If repo specified, print archives for only that repo
  if [[ $# -gt 0 ]]; then
    while [[ $# -gt 0 ]]; do
      repoData="$(select_repo_by_name "$1")"
      list_repo_archives_helper "${repoData}"

      printf "\n"
      shift
    done
  else
    # If repo not specified, print archives for all repos
    print_script_message "Listing repo archives."

    local whileCounter=0
    while [[ "${whileCounter}" -lt "${KBH_REPOS_NUMBER}" ]]; do
      repoData="$(select_repo_by_index "${whileCounter}")"
      list_repo_archives_helper "${repoData}"

      whileCounter="$(( whileCounter + 1 ))"
      printf "\n"
    done
  fi
}

backup_helper()
{
  local -r repoName="$(jaq -r .name <<< "$1")"
  local -r repoPath="$(jaq -r .path <<< "$1")"
  local -r passwordCommand="$(jaq -r .password_command <<< "$1")"
  local -r archiveName="$(${KBH_ARCHIVE_NAME})"

  local backupPaths
  mapfile -t backupPaths <<< "$(jaq -r .backup_paths[] <<< "${KBH_SETTINGS_JSON}")"
  readonly backupPaths

  local excludePaths
  mapfile -t excludePaths <<< "$(jaq -r .exclude_paths[] <<< "${KBH_SETTINGS_JSON}")"
  readonly excludePaths

  local borg_exclude_command_slice=" "

  if [[ "${excludePaths}" != "null" ]] && [[ "${#excludePaths[@]}" -gt 0  ]]; then
    for path in "${excludePaths[@]}"; do
      borg_exclude_command_slice="${borg_exclude_command_slice}-e ${path}"
    done
  fi

  local repoCompression
  repoCompression="$(jaq -r .compression <<< "$1")"
  if [[ "${repoCompression}" == "null" ]]; then
    # Use default borg compression
    repoCompression="lz4"
  fi
  readonly repoCompression

  print_script_message "Backing up files to ${repoName}..."
  printf "Running command: borg create --compression %s --stats --progress %s::%s %s%s\n\n" "${repoCompression}" "${repoPath}" "${archiveName}" "${backupPaths[*]}" "${borg_exclude_command_slice}"

  # Must use shellcheck disable as borg explicitly wants regexp
  # shellcheck disable=SC2086
  BORG_PASSPHRASE="$(${passwordCommand})" borg create --compression "${repoCompression}" --stats --progress "${repoPath}"::"${archiveName}" "${backupPaths[@]}" ${borg_exclude_command_slice}
}

# $*: repos (optional)
backup()
{
  shift
  local repoData=""

  # If repo specified, backup only to that repo
  if [[ $# -gt 0 ]]; then
    while [[ $# -gt 0 ]]; do
      repoData="$(select_repo_by_name "$1")"
      backup_helper "${repoData}"

      printf "\n"
      shift
    done
  else
    # If repo not specified, backup to all repos
    print_script_message "Backing up to all repos."

    local whileCounter=0
    while [[ "${whileCounter}" -lt "${KBH_REPOS_NUMBER}" ]]; do
      repoData="$(select_repo_by_index "${whileCounter}")"
      backup_helper "${repoData}"

      whileCounter="$(( whileCounter + 1 ))"
      printf "\n"
    done
  fi
}

prune_repos_helper()
{
  local -r repoName="$(jaq -r .name <<< "$1")"
  local -r repoPath="$(jaq -r .path <<< "$1")"
  local -r passwordCommand="$(jaq -r .password_command <<< "$1")"

  local -r pruneCommandSetting="$(jaq -r .prune_command <<< "${KBH_SETTINGS_JSON}")"
  local pruneCommandOfficial=""

  # If prune command not specified, use default as stated in man page
  if [[ "${pruneCommandSetting}" == "null" ]]; then
    pruneCommandOfficial="borg prune -v --list --keep-daily=30 --keep-weekly=-1"
  elif [[ -z "${pruneCommandOfficial}" ]]; then
    pruneCommandOfficial="${pruneCommandSetting}"
  else
    exit_script_on_failure "Could not set prune command"
  fi

  readonly pruneCommandOfficial

  # Do dry run first and prompt for confirmation
  print_script_message "Doing DRY RUN for pruning archives on ${repoName}..."
  printf "Running command: %s --dry-run %s\n" "${pruneCommandOfficial}" "${repoPath}"
  BORG_PASSPHRASE="$(${passwordCommand})" ${pruneCommandOfficial} --dry-run "${repoPath}"
  
  printf "\n%sApprove execution of prune command? " "${WHITE}"
  read -r -p "[y/N]${NC} " response

  case "${response}" in
    [yY][eE][sS]|[yY])
      printf "\n"
      print_script_message "Pruning archives on ${repoName}..."
      printf "Running command: %s %s\n" "${pruneCommandOfficial}" "${repoPath}"
      BORG_PASSPHRASE="$(${passwordCommand})" ${pruneCommandOfficial} "${repoPath}"

      printf "\n"
      print_script_message "Compacting archives on ${repoName}..."
      printf "Running command: borg compact %s\n" "${repoPath}"
      BORG_PASSPHRASE="$(${passwordCommand})" borg compact "${repoPath}"
      ;;
    *)
      printf "Cancelling prune command\n"
      ;;
  esac
}

# $*: repos (optional)
prune_repos()
{
  shift
  local repoData=""

  # If repo specified, prune only that repo
  if [[ $# -gt 0 ]]; then
    while [[ $# -gt 0 ]]; do
      repoData="$(select_repo_by_name "$1")"
      prune_repos_helper "${repoData}"

      printf "\n"
      shift
    done
  else
    # If repo not specified, prune all repos
    print_script_message "Pruning all repos."

    local whileCounter=0
    while [[ "${whileCounter}" -lt "${KBH_REPOS_NUMBER}" ]]; do
      repoData="$(select_repo_by_index "${whileCounter}")"
      prune_repos_helper "${repoData}"

      whileCounter="$(( whileCounter + 1 ))"
      printf "\n"
    done
  fi
}

# no arguments
init_repos()
{
  local repoName=""
  local repoPath=""
  local passwordCommand=""
  local whileCounter=0

  print_script_message "Creating borg repositories..."

  while [[ "${whileCounter}" -lt "${KBH_REPOS_NUMBER}" ]]; do
    repoName="$(jaq -r .repos[${whileCounter}].name <<< "${KBH_SETTINGS_JSON}")"
    repoPath="$(jaq -r .repos[${whileCounter}].path <<< "${KBH_SETTINGS_JSON}")"
    passwordCommand="$(jaq -r .repos[${whileCounter}].password_command <<< "${KBH_SETTINGS_JSON}")"
    
    # Test if repo already present; if so, skip
    if BORG_PASSPHRASE="$(${passwordCommand})" borg info "${repoPath}" &> /dev/null; then
      printf "%sBorg repo already present at %s%s%s for repo %s%s%s\n" "${YELLOW}" "${NC}" "${repoPath}" "${YELLOW}" "${WHITE}" "${repoName}" "${NC}"
      whileCounter="$(( whileCounter + 1 ))"

      continue
    fi

    print_script_message "Creating borg repo for ${repoName}..."
    borg init --encryption repokey-blake2 "${repoPath}"

    print_script_message "Printing borg key from ${repoName}:"
    BORG_PASSPHRASE="$(${passwordCommand})" borg key export "${repoPath}"
    
    print_script_message "--------------------------------------------"

    whileCounter="$(( whileCounter + 1 ))"
    printf "\n"
  done
}

# $1: repo
# $2: archive (optional)
mount_archive()
{
  shift

  if [[ -z "${1:-}" ]]; then
    exit_script_on_failure "No repository specified, e.g. ${WHITE}kbh mount home-server${NC}"
  fi

  local -r mountPoint="$(get_mount_point)"

  # Ensure directory $mountPoint exists
  if [[ ! -d "${mountPoint}" ]]; then
    exit_script_on_failure "Mount point ${WHITE}${mountPoint}${NC} directory does not exist"
  fi

  # Ensure $mountPoint does not have something mounted on it already
  if mountpoint "${mountPoint}" &> /dev/null; then
    exit_script_on_failure "${mountPoint} already has something mounted on it"
  fi

  # Find repo data for specified repo
  local -r repoData="$(select_repo_by_name "$1")"

  # If no archive specified, mount repo
  if [[ -z "${2:-}" ]]; then
    BORG_PASSCOMMAND="$(jaq -r .password_command <<< "${repoData}")" borg mount "$(jaq -r .path <<< "${repoData}")" "${mountPoint}"
  else
    # If archive specified, mount archive
    BORG_PASSCOMMAND="$(jaq -r .password_command <<< "${repoData}")" borg mount "$(jaq -r .path <<< "${repoData}")"::"$2" "${mountPoint}"
  fi
}

# Print config and overall repo info
print_status()
{
  # Print base settings
  mapfile -t backupPaths <<< "$(jaq -r .backup_paths[] <<< "${KBH_SETTINGS_JSON}")"
  local -r sshKeyPasswordCopyCommand="$(jaq -r .ssh_key_password_copy_command <<< "${KBH_SETTINGS_JSON}")"
  local -r borgArchiveNameCommand="$(jaq -r .borg_archive_name_command <<< "${KBH_SETTINGS_JSON}")"
  local -r mountPoint="$(get_mount_point)"

  # Custom printf line cuz we fancy
  printf "%s==>%s KBH settings %s(%s%s%s)\n" "${GREEN}" "${WHITE}" "${NC}" "${GREY}" "${CONFIG_FILE}" "${NC}"

  printf "Backup paths:          %s%s%s\n" "${WHITE}" "${backupPaths[*]}" "${NC}"
  printf "Borg archive name cmd: %s%s (e.g. %s)%s\n" "${WHITE}" "${borgArchiveNameCommand}" "$(${borgArchiveNameCommand})" "${NC}"
  printf "SSH key pwd copy cmd:  %s%s%s\n" "${WHITE}" "${sshKeyPasswordCopyCommand}" "${NC}"
  printf "Mount point:           %s%s%s\n\n" "${WHITE}" "${mountPoint}" "${NC}"

  # Print per-repo information
  local whileCounter=0
  local repoName=""
  local repoPath=""
  local passwordCommand=""
  local repoCompression=""
  local borgInfoOutput=""
  local lastArchive=""

  print_script_message "Repository information"

  while [[ "${whileCounter}" -lt "${KBH_REPOS_NUMBER}" ]]; do
    repoName="$(jaq -r .repos[${whileCounter}].name <<< "${KBH_SETTINGS_JSON}")"
    repoPath="$(jaq -r .repos[${whileCounter}].path <<< "${KBH_SETTINGS_JSON}")"
    passwordCommand="$(jaq -r .repos[${whileCounter}].password_command <<< "${KBH_SETTINGS_JSON}")"
    repoCompression="$(jaq -r .repos[${whileCounter}].compression <<< "${KBH_SETTINGS_JSON}")"

    if [[ "${repoCompression}" == "null" ]]; then
      # Use default borg compression
      repoCompression="lz4 (default)"
    fi

    # If we can connect to repo, print more info (repo size and last backup info)
    if ! BORG_PASSPHRASE="$(${passwordCommand})" timeout 8s borg info "${repoPath}" &> /dev/null; then
      printf "%s%s%s (could not connect or find repo)%s\n" "${WHITE}" "${repoName}" "${YELLOW}" "${NC}"
    else
      borgInfoOutput="$(BORG_PASSPHRASE="$(${passwordCommand})" borg info "${repoPath}")"
      borgListOutput="$(BORG_PASSPHRASE="$(${passwordCommand})" borg list "${repoPath}")"
      lastArchive="$(printf "%s" "${borgListOutput}" | tail -n1 | tr -s " ")"

      printf "%s%s%s (%s %s, %s archives)\n" "${WHITE}" "${repoName}" "${NC}" "$(printf "%s" "${borgInfoOutput}" | grep 'All archives' | awk '{print $7}')" "$(printf "%s" "${borgInfoOutput}" | grep 'All archives' | awk '{print $8}')" "$(printf "%s\n" "${borgListOutput}" | wc -l)"
      printf "  Last backup: %s\n" "${lastArchive% *}"
    fi

    printf "  Path:        %s\n" "${repoPath}"
    printf "  Pwd cmd:     %s\n" "${passwordCommand}"
    printf "  Compression: %s\n" "${repoCompression}"

    printf "\n"

    whileCounter="$(( whileCounter + 1 ))"
  done
}

delete()
{
  # NYI
  true
}

# -----------------------------------------
# ---------------- Script -----------------
# -----------------------------------------

if [[ "$(whoami)" = "root" ]]; then
  exit_script_on_failure "This script should NOT be run as root (or sudo)!"
fi

# Check user input
[[ -z "${1:-}" ]] && print_help

check_ssh_key

case "$1" in
  help|-h|--help) print_help "$@" ;;
  prune) prune_repos "$@" ;;
  info) repo_info "$@" ;;
  init) init_repos ;;
  list) list_repo_archives "$@" ;;
  mount) mount_archive "$@" ;;
  backup) backup "$@" ;;
  delete) delete "$@" ;;
  status) print_status ;;
  *) exit_script_on_failure "Invalid or missing subcommand. Use \`kbh help\` to list subcommands, or read the man page with \`man kbh\`."
esac

