#!/bin/bash
#
# Restores a MariaDB Galera Cluster node from a backup produced by
# backup_galera_cluster_full.sh.
#
# The script presents an interactive selection of available full backups and
# their associated incremental backups, allowing the operator to choose exactly
# which point in time to restore to. This is intentional: the latest backup is
# not necessarily a good one.
#
# The existing data directory is moved aside before the restore so it can be
# recovered manually if needed.
#
# After the restore completes, the cluster must be re-bootstrapped from this
# node. See the documentation for the full procedure.
#
# Usage: restore_galera_cluster_full.sh
#   Must be run as root. MariaDB must be stopped on ALL nodes before running.

# --- Configuration -----------------------------------------------------------

readonly BACK_DIR='/mariadb_backup'
readonly MYSQL_DATA_DIR='/var/lib/mysql'
readonly BACK_CMD='mariadb-backup'
readonly GZIP_CMD='gzip'
readonly STREAM_CMD='mbstream'

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

#######################################
# Print an error message to STDERR with a timestamp.
# Arguments:
#   Message string(s) to print.
# Outputs:
#   Writes timestamped message to STDERR.
#######################################
err() {
  echo "[$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*" >&2
}

#######################################
# Verify that all required external commands are available in PATH.
# Arguments:
#   None
# Returns:
#   Exits with code 1 if any required command is not found.
#######################################
check_dependencies() {
  local -r required_cmds=('mariadb-backup' 'mbstream' 'gzip')
  local cmd

  for cmd in "${required_cmds[@]}"; do
    if ! command -v "${cmd}" >/dev/null 2>&1; then
      err "Required command not found: ${cmd}"
      err "Ensure it is installed and available in PATH."
      exit 1
    fi
  done
}

#######################################
# Verify that the script is run as root.
# Arguments:
#   None
# Returns:
#   Exits with code 1 if not running as root.
#######################################
check_root() {
  if ((EUID != 0)); then
    err 'This script must be run as root.'
    exit 1
  fi
}

#######################################
# Verify that MariaDB is not running. A restore must not be performed while
# the data directory is in active use.
# Arguments:
#   None
# Returns:
#   Exits with code 1 if MariaDB is running.
#######################################
check_mariadb_stopped() {
  if systemctl is-active --quiet mariadb.service; then
    err 'MariaDB is still running. Stop it on ALL nodes before restoring:'
    err '  systemctl stop mariadb.service'
    exit 1
  fi
}

#######################################
# Present a numbered list of options and prompt the operator to select one.
# Arguments:
#   prompt:       The question to display above the list.
#   option_label: Label shown before each item (e.g. "Full backup").
#   items:        All remaining arguments are treated as selectable items.
# Outputs:
#   Writes the selected item to STDOUT.
# Returns:
#   Exits with code 1 if no valid selection is made after 3 attempts,
#   or if the operator chooses to abort.
#######################################
select_item() {
  local -r prompt="$1"
  local -r option_label="$2"
  shift 2
  local -r items=("$@")
  local -r count="${#items[@]}"

  echo >&2
  echo "${prompt}" >&2
  echo >&2

  local i
  for ((i = 0; i < count; i++)); do
    printf '  [%d] %s: %s\n' "$((i + 1))" "${option_label}" "${items[${i}]}" >&2
  done
  echo '  [0] Abort restore' >&2
  echo >&2

  local attempts=0
  local choice
  while ((attempts < 3)); do
    read -r -p "Enter selection [0-${count}]: " choice </dev/tty

    if [[ "${choice}" == '0' ]]; then
      echo 'Restore aborted.' >&2
      exit 0
    fi

    if [[ "${choice}" =~ ^[0-9]+$ ]] &&
      ((choice >= 1 && choice <= count)); then
      # Only the selected value goes to STDOUT for capture by the caller.
      echo "${items[$((choice - 1))]}"
      return 0
    fi

    err "Invalid selection: ${choice}. Please enter a number between 0 and ${count}."
    ((attempts++))
  done

  err 'Too many invalid attempts. Aborting.'
  exit 1
}

#######################################
# Prompt the operator to select a full backup from those available.
# Globals:
#   BACK_DIR
# Arguments:
#   None
# Outputs:
#   Writes the full path of the selected base backup directory to STDOUT.
# Returns:
#   Exits with code 1 if no full backups are found.
#######################################
select_full_backup() {
  local -a full_backups
  readarray -t full_backups < <(find "${BACK_DIR}/base" \
    -mindepth 1 -maxdepth 1 -type d |
    sort -r)

  if ((${#full_backups[@]} == 0)); then
    err "No full backups found in ${BACK_DIR}/base."
    exit 1
  fi

  select_item \
    'Available full backups (newest first):' \
    'Full backup' \
    "${full_backups[@]}"
}

#######################################
# Prompt the operator to select how many incrementals to apply on top of the
# chosen full backup. The operator can choose any incremental in the chain,
# or skip incrementals entirely.
# Globals:
#   BACK_DIR
# Arguments:
#   full_backup_name: directory name (not full path) of the chosen full backup.
# Outputs:
#   Writes a newline-separated list of selected incremental paths to STDOUT,
#   in chronological order. Writes nothing if no incrementals are to be applied.
#######################################
select_incrementals() {
  local -r full_backup_name="$1"
  local -r incr_dir="${BACK_DIR}/incr/${full_backup_name}"

  if [[ ! -d "${incr_dir}" ]]; then
    echo "No incremental backups found for this full backup." >&2
    return 0
  fi

  local -a incrementals
  readarray -t incrementals < <(find "${incr_dir}" \
    -mindepth 1 -maxdepth 1 -type d |
    sort)

  if ((${#incrementals[@]} == 0)); then
    echo "No incremental backups found for this full backup." >&2
    return 0
  fi

  echo >&2
  echo "Available incremental backups for selected full (oldest first):" >&2
  echo >&2
  echo "  [0] Do not apply any incrementals (restore full backup only)" >&2

  local i
  for ((i = 0; i < ${#incrementals[@]}; i++)); do
    printf '  [%d] Apply up to: %s\n' \
      "$((i + 1))" "$(basename "${incrementals[${i}]}")" >&2
  done
  echo >&2

  local attempts=0
  local choice
  local count="${#incrementals[@]}"
  while ((attempts < 3)); do
    read -r -p "Apply how many incrementals? [0-${count}]: " choice </dev/tty

    if [[ "${choice}" =~ ^[0-9]+$ ]] &&
      ((choice >= 0 && choice <= count)); then
      # Print only the selected slice of incrementals to STDOUT.
      local j
      for ((j = 0; j < choice; j++)); do
        echo "${incrementals[${j}]}"
      done
      return 0
    fi

    err "Invalid selection: ${choice}. Please enter a number between 0 and ${count}."
    ((attempts++))
  done

  err 'Too many invalid attempts. Aborting.'
  exit 1
}

#######################################
# Decompress and unstream a single backup stream into a target directory.
# Arguments:
#   source_file: path to the backup.stream.gz file.
#   target_dir:  directory to extract the backup into.
# Returns:
#   Exits with code 1 if decompression or streaming fails.
#######################################
decompress_backup() {
  local -r source_file="$1"
  local -r target_dir="$2"

  echo "Decompressing: ${source_file} -> ${target_dir}"
  mkdir -p "${target_dir}"

  "${GZIP_CMD}" -dc "${source_file}" |
    "${STREAM_CMD}" -x -C "${target_dir}"

  local -r pipe_status=("${PIPESTATUS[@]}")
  if ((pipe_status[0] != 0 || pipe_status[1] != 0)); then
    err "Decompression failed for ${source_file}"
    err "  gzip exit code:     ${pipe_status[0]}"
    err "  mbstream exit code: ${pipe_status[1]}"
    exit 1
  fi
}

#######################################
# Prepare the full backup and apply the selected incremental backups on top
# of it in chronological order. Each step runs --prepare, which is correct
# for MariaDB 11.4. --apply-log-only is not required or supported.
# Globals:
#   BACK_CMD
# Arguments:
#   full_backup_dir: path to the decompressed full backup working directory.
#   incrementals:    newline-separated list of incremental paths to apply,
#                    in chronological order. May be empty.
# Returns:
#   Exits with code 1 if any prepare step fails.
#######################################
prepare_backup() {
  local -r full_backup_dir="$1"
  local -r incrementals="$2"

  # Prepare the full backup, leaving it open for incremental application.
  echo 'Preparing full backup...'
  if ! "${BACK_CMD}" --prepare \
    --target-dir="${full_backup_dir}"; then
    err 'Failed to prepare full backup.'
    exit 1
  fi

  # Apply each selected incremental in order.
  if [[ -n "${incrementals}" ]]; then
    local incr
    while IFS= read -r incr; do
      local incr_work_dir
      incr_work_dir="${full_backup_dir}/incr_$(basename "${incr}")"

      decompress_backup "${incr}/backup.stream.gz" "${incr_work_dir}"

      echo "Applying incremental: $(basename "${incr}")"
      if ! "${BACK_CMD}" --prepare \
        --target-dir="${full_backup_dir}" \
        --incremental-dir="${incr_work_dir}"; then
        err "Failed to apply incremental backup: ${incr}"
        exit 1
      fi

      rm -rf "${incr_work_dir}"
    done <<<"${incrementals}"
  else
    echo 'No incrementals selected. Restoring full backup only.'
  fi
}

#######################################
# Move the existing MariaDB data directory aside and copy the prepared backup
# into its place.
# Globals:
#   MYSQL_DATA_DIR, BACK_CMD
# Arguments:
#   full_backup_dir: path to the prepared backup working directory.
# Returns:
#   Exits with code 1 if the data directory move or copy-back fails.
#######################################
restore_data_dir() {
  local -r full_backup_dir="$1"
  local -r archive_dir="${MYSQL_DATA_DIR}_$(date +%F_%H-%M-%S)"

  echo "Moving existing data directory to: ${archive_dir}"
  if ! mv "${MYSQL_DATA_DIR}" "${archive_dir}"; then
    err "Failed to move ${MYSQL_DATA_DIR} to ${archive_dir}."
    exit 1
  fi

  echo "Copying backup into: ${MYSQL_DATA_DIR}"
  if ! "${BACK_CMD}" --copy-back \
    --target-dir="${full_backup_dir}"; then
    err 'Failed to copy backup into data directory.'
    err "The original data directory has been preserved at: ${archive_dir}"
    exit 1
  fi

  echo "Restoring ownership of ${MYSQL_DATA_DIR} to mysql:mysql"
  if ! chown -R mysql:mysql "${MYSQL_DATA_DIR}"; then
    err "Failed to set ownership on ${MYSQL_DATA_DIR}."
    exit 1
  fi
}

#######################################
# Main entry point.
#######################################
main() {
  local -r start_time="$(date +%s)"
  local work_dir
  work_dir="$(mktemp -d /tmp/mariadb_restore.XXXXXX)"
  readonly work_dir

  # Ensure the temporary working directory is always cleaned up on exit.
  trap 'echo "Cleaning up temporary files."; rm -rf "${work_dir}"' EXIT
  trap 'err "Interrupted."; exit 1' INT TERM

  check_root
  check_dependencies
  check_mariadb_stopped

  echo '============================================================'
  echo ' MariaDB Galera Cluster Restore'
  echo '============================================================'
  echo "Backup source:  ${BACK_DIR}"
  echo "Data directory: ${MYSQL_DATA_DIR}"
  echo "Working dir:    ${work_dir}"

  # Step 1: select full backup.
  local selected_full
  selected_full="$(select_full_backup)"
  local -r selected_full_name="$(basename "${selected_full}")"

  # Step 2: select incrementals to apply.
  local selected_incrementals
  selected_incrementals="$(select_incrementals "${selected_full_name}")"

  # Summarise the restore plan and ask for explicit confirmation.
  local incr_count=0
  if [[ -n "${selected_incrementals}" ]]; then
    incr_count="$(echo "${selected_incrementals}" | wc -l)"
  fi

  echo
  echo '--- Restore plan -----------------------------------------------'
  echo "  Full backup:            ${selected_full_name}"
  echo "  Incrementals to apply:  ${incr_count}"
  if [[ -n "${selected_incrementals}" ]]; then
    local incr
    while IFS= read -r incr; do
      echo "    - $(basename "${incr}")"
    done <<<"${selected_incrementals}"
  fi
  echo "  Data dir to replace:    ${MYSQL_DATA_DIR}"
  echo '----------------------------------------------------------------'
  echo

  read -r -p "Proceed with restore? This will replace ${MYSQL_DATA_DIR}. [yes/N] " confirm </dev/tty
  if [[ "${confirm}" != 'yes' ]]; then
    echo 'Restore cancelled.'
    exit 0
  fi

  # Step 3: decompress the selected full backup into the working directory.
  local -r full_work_dir="${work_dir}/full"
  decompress_backup "${selected_full}/backup.stream.gz" "${full_work_dir}"

  # Step 4: prepare the full backup and apply selected incrementals.
  prepare_backup "${full_work_dir}" "${selected_incrementals}"

  # Step 5: swap the data directory and copy the backup in.
  restore_data_dir "${full_work_dir}"

  local -r elapsed=$(($(date +%s) - start_time))
  echo
  echo '============================================================'
  echo " Restore completed in ${elapsed} seconds at $(date)."
  echo '============================================================'
  echo
  echo 'Next steps:'
  echo '  1. Bootstrap the cluster from THIS node:'
  echo '       galera_new_cluster'
  echo '  2. Verify the node is healthy:'
  echo '       journalctl -n 100 -f -u mariadb.service'
  echo '  3. Start MariaDB on all other nodes one at a time:'
  echo '       systemctl start mariadb.service'
  echo '  4. Verify each node shows wsrep_local_state_comment = Synced'
  echo '  5. Restart application services.'
}

main "$@"