{% raw %}#!/usr/bin/env bash
set -euo pipefail

# ===== CONFIG (filled by Ansible) =====
KEEP_LOCAL=30
KEEP_REMOTE=10
KEYFILE="/root/.ssh/backup"
HOST_DIR="backup"
DATASET="{% endraw %}{{ backup_zfs_dataset }}{% raw %}"
BACKUP_EXCLUDE_PROP="{% endraw %}{{ backup_zfs_exclude_property | default('com.pursotin:backup') }}{% raw %}"
USER="{% endraw %}{{ backup_ssh_user }}{% raw %}"
HOST="{% endraw %}{{ backup_ssh_host }}{% raw %}"
EMAIL_TO=root
# =====================================

# ----- Globals for notification -----
STARTED_AT="$(date '+%Y-%m-%dT%H:%M:%S%z')"
FINISHED_AT=""
MESSAGE_LOG=""
BACKUP_NAME=""          # e.g. zroot@2025-09-30-12-00-00-0300.enc
BACKUP_SIZE_BYTES=""    # numeric bytes

# ----- Helpers -----
die() { echo "Error: $*" >&2; exit 1; }

require_cmds() {
  local cmds=("$@")
  for c in "${cmds[@]}"; do command -v "$c" >/dev/null 2>&1 || die "Missing command: $c"; done
}

log_note() {
  # Echo to console and append to message buffer
  local msg="$1"
  echo "$msg"
  MESSAGE_LOG+="$msg"$'\n'
}

timestamp() {
  # Replace '+' with '-' so timezone is filename-safe and lexicographically sortable within TZ.
  date +%Y-%m-%d-%H-%M-%S%z | tr '+' '-'
}

humanize_bytes() {
  local bytes="$1"
  if [[ "${bytes}" =~ ^[0-9]+$ ]]; then
    awk -v b="${bytes}" '
      BEGIN {
        unit_count = split("B KiB MiB GiB TiB PiB EiB", units, " ")
        i = 1
        while (b >= 1024 && i < unit_count) {
          b = b / 1024
          i++
        }
        if (i == 1) {
          printf "%.0f %s\n", b, units[i]
        } else {
          printf "%.2f %s\n", b, units[i]
        }
      }
    '
  else
    echo "unknown"
  fi
}

latest_snapshot_for_dataset() {
  # Latest snapshot on the TOP dataset only (newest first). Returns e.g. zroot@2025-09-30-...
  zfs list -t snapshot -o name -S creation "$DATASET" 2>/dev/null \
    | awk 'NR==2{print; exit}'
}

list_top_snapshots_newest_first() {
  # Only list snapshots on the top dataset (not children), newest first.
  zfs list -t snapshot -o name -S creation "$DATASET" 2>/dev/null \
    | awk 'NR>1{print $1}'
}

list_excluded_datasets() {
  # List datasets (including inherited values) where backup property is explicitly false.
  zfs get -r -H -o name,value "${BACKUP_EXCLUDE_PROP}" "${DATASET}" 2>/dev/null \
    | awk 'tolower($2)=="false"{print $1}'
}

list_remote_backups_sorted() {
  # We store files under ${HOST_DIR}/<dataset>@<timestamp>.enc
  echo "ls ${HOST_DIR}" \
    | sftp -q -i "${KEYFILE}" "${USER}@${HOST}" 2>/dev/null \
    | tail -n +2 | sort
}

# Upload a local file to the storage box path "${HOST_DIR}/<remote_name>"
sftp_put() {
  local local_file="$1"
  local remote_name="$2"  # just the filename, no directory
  sftp -i "${KEYFILE}" "${USER}@${HOST}" <<EOF
put ${local_file} ${HOST_DIR}/${remote_name}
EOF
}

# ----- Subcommands -----
cmd_init_remote() {
  echo ""
  echo "[init-remote]"
  require_cmds sftp

  echo "Checking for remote directory: ${HOST_DIR}"
  if echo "ls ${HOST_DIR}" | sftp -i "${KEYFILE}" "${USER}@${HOST}" >/dev/null 2>&1; then
    echo "Remote directory already exists: ${HOST_DIR}"
    log_note "init-remote: verified remote directory '${HOST_DIR}'"
    return 0
  fi

  echo "Creating remote directory: ${HOST_DIR}"
  set +e
  sftp -i "${KEYFILE}" "${USER}@${HOST}" <<EOF
mkdir ${HOST_DIR}
EOF
  rc=$?
  set -e
  if [[ $rc -ne 0 ]]; then
    die "Failed to create remote directory ${HOST_DIR}"
  fi

  # Re-check
  if echo "ls ${HOST_DIR}" | sftp -q -i "${KEYFILE}" "${USER}@${HOST}" >/dev/null 2>&1; then
    echo "Remote directory created: ${HOST_DIR}"
    log_note "init-remote: created remote directory '${HOST_DIR}'"
  else
    die "Remote directory ${HOST_DIR} not found after creation"
  fi
}

cmd_snapshot() {
  echo ""
  echo "[snapshot]"
  require_cmds zfs date tr
  local ts snap
  ts="$(timestamp)"
  snap="${DATASET}@${ts}"
  echo "Taking recursive snapshot \"${snap}\""
  (set -x; zfs snapshot -r "${snap}")
  echo "Recursive snapshot \"${snap}\" created"
  log_note "snapshot: created recursive snapshot '${snap}'"
}

cmd_send_to_remote() {
  echo ""
  echo "[send-to-remote]"
  require_cmds zfs age sftp mktemp stat awk sort

  local snap base target tmp size_bytes ds parent skip
  local -a send_args excluded_raw excluded
  snap="$(latest_snapshot_for_dataset)"
  [[ -n "${snap}" ]] || die "No snapshot found to send. Run 'snapshot' first or ensure the dataset has snapshots."

  base="$(basename "${snap}")"       # e.g., zroot@2025-09-30-12-00-00-0300
  target="${base}.enc"
  tmp="$(mktemp -t backup_send.XXXXXX)"
  trap 'rm -f "${tmp}"' EXIT

  send_args=(-Rvc)
  mapfile -t excluded_raw < <(list_excluded_datasets | sort)
  excluded=()
  for ds in "${excluded_raw[@]}"; do
    skip=0
    [[ -n "${ds}" ]] || continue
    for parent in "${excluded[@]}"; do
      if [[ "${ds}" == "${parent}" || "${ds}" == "${parent}/"* ]]; then
        skip=1
        break
      fi
    done
    [[ "${skip}" -eq 1 ]] || excluded+=("${ds}")
  done

  if [[ "${#excluded[@]}" -gt 0 ]]; then
    echo "Excluding datasets where ${BACKUP_EXCLUDE_PROP}=false:"
    for ds in "${excluded[@]}"; do
      if [[ "${ds}" == "${DATASET}" ]]; then
        die "Top dataset ${DATASET} is excluded via ${BACKUP_EXCLUDE_PROP}=false; refusing to create an empty backup stream."
      fi
      echo "  - ${ds}"
      send_args+=(-X "${ds}")
    done
    log_note "send-to-remote: excluded ${#excluded[@]} dataset(s) using ${BACKUP_EXCLUDE_PROP}=false"
  fi

  echo "Creating encrypted replication stream to temp file: ${tmp}"
  # -R: recursive hierarchy, -v: progress to stderr, -c: send compressed (keeps on-disk compression)
  (set -x; zfs send "${send_args[@]}" "${snap}" | age -e -i "${KEYFILE}" > "${tmp}")

  size_bytes="$(stat -f %z "${tmp}" 2>/dev/null || stat -c %s "${tmp}" 2>/dev/null || echo "unknown")"
  echo "Local stream size: ${size_bytes} bytes"

  echo "Uploading via SFTP to ${HOST}:${HOST_DIR}/${target}"
  if sftp_put "${tmp}" "${target}"; then
    echo "Upload complete"
  else
    die "SFTP upload failed (does the remote directory '${HOST_DIR}' exist? Run 'init-remote')"
  fi

  # Set globals for notify()
  BACKUP_NAME="${target}"
  BACKUP_SIZE_BYTES="${size_bytes}"

  log_note "send-to-remote: uploaded '${snap}' as '${target}' (${size_bytes} bytes) to ${HOST}:${HOST_DIR}"
  rm -f "${tmp}"
  trap - EXIT
}

cmd_prune_remote() {
  echo ""
  echo "[prune-remote]"
  require_cmds sftp awk sort wc

  echo "Fetching remote backup listing from sftp://${HOST}/${HOST_DIR}"
  BACKUPS="$(list_remote_backups_sorted || true)"
  mapfile -t BACKUP_ARR < <(printf "%s\n" "${BACKUPS}")
  local count="${#BACKUP_ARR[@]}"

  if [[ "${count}" -le "${KEEP_REMOTE}" ]]; then
    echo "Remote backups (${count}) <= KEEP_REMOTE (${KEEP_REMOTE}); nothing to prune."
    log_note "prune-remote: kept ${count} (<= ${KEEP_REMOTE}); no deletions"
    return 0
  fi

  echo "Pruning remote backups, keeping latest ${KEEP_REMOTE} (will delete $(("${count}" - "${KEEP_REMOTE}")))"
  local to_delete_count=$((count - KEEP_REMOTE))
  local deleted=0
  for ((i=0; i<to_delete_count; i++)); do
    b="${BACKUP_ARR[$i]}"
    [[ -n "${b}" ]] || continue
    echo "Deleting remote: ${b}"
    set +e
    echo "rm ${b}" | sftp -i "${KEYFILE}" "${USER}@${HOST}"
    rc=$?
    set -e
    if [[ $rc -ne 0 ]]; then
      die "Failed to delete remote file: ${b}"
    fi
    deleted=$((deleted+1))
  done
  echo "Remote prune done. Deleted: ${deleted}"
  log_note "prune-remote: deleted ${deleted}, kept ${KEEP_REMOTE}"
}

cmd_prune_local() {
  echo ""
  echo "[prune-local]"
  require_cmds zfs awk

  echo "Pruning local snapshots on dataset: ${DATASET}; keeping latest ${KEEP_LOCAL}"
  mapfile -t snaps < <(list_top_snapshots_newest_first)
  local total="${#snaps[@]}"
  if [[ "${total}" -le "${KEEP_LOCAL}" ]]; then
    echo "Local snapshots (${total}) <= KEEP_LOCAL (${KEEP_LOCAL}); nothing to prune."
    log_note "prune-local: kept ${total} (<= ${KEEP_LOCAL}); no deletions"
    return 0
  fi

  local deleted=0
  # List is newest-first; skip first KEEP_LOCAL and destroy the rest RECURSIVELY across the tree
  for ((i=KEEP_LOCAL; i<total; i++)); do
    s="${snaps[$i]}"   # e.g., zroot@2025-09-30-...
    echo "Destroying recursive snapshot: ${s}"
    (set -x; zfs destroy -r "${s}")
    deleted=$((deleted+1))
  done
  echo "Local prune done. Deleted: ${deleted}"
  log_note "prune-local: deleted ${deleted}, kept ${KEEP_LOCAL}"
}

cmd_notify() {
  echo ""
  echo "[notify]"
  require_cmds mail date awk
  local backup_size_human
  FINISHED_AT="$(date '+%Y-%m-%dT%H:%M:%S%z')"

  # Derive snapshot name if not set yet (best-effort)
  if [[ -z "${BACKUP_NAME}" ]]; then
    # Try to infer from latest snapshot
    latest="$(latest_snapshot_for_dataset || true)"
    if [[ -n "${latest}" ]]; then
      BACKUP_NAME="$(basename "${latest}").enc"
    else
      BACKUP_NAME="unknown"
    fi
  fi

  backup_size_human="$(humanize_bytes "${BACKUP_SIZE_BYTES:-unknown}")"

  # Compose subject and body
  local subject="Backup completed: ${BACKUP_NAME}"
  local body=""
  body+="Backup run summary"$'\n'
  body+="Started:  ${STARTED_AT}"$'\n'
  body+="Finished: ${FINISHED_AT}"$'\n'
  body+="Dataset:  ${DATASET}"$'\n'
  body+="Remote:   ${USER}@${HOST}:${HOST_DIR}"$'\n'
  body+="Name:     ${BACKUP_NAME}"$'\n'
  body+="Size:     ${backup_size_human} (${BACKUP_SIZE_BYTES:-unknown} bytes)"$'\n'
  body+=$'\n'
  body+="Steps:"$'\n'
  body+="${MESSAGE_LOG:-<no steps recorded>}"$'\n'

  echo "Sending notification email..."
  printf "%s\n" "$body" | mail -s "$subject" "${EMAIL_TO}"
  log_note "notify: email sent to ${EMAIL_TO}"
}

usage() {
  cat <<EOF
Usage: $(basename "$0") [subcommand [subcommand ...]]

Subcommands (executed in order):
  init-remote      Create the remote backup directory (\$HOST_DIR). Succeeds if it already exists.
  snapshot         Create a new **recursive** ZFS snapshot for \$DATASET
  send-to-remote   Create encrypted replication stream to a temp file, then upload via SFTP
  prune-remote     Keep latest \$KEEP_REMOTE backups on remote (default ${KEEP_REMOTE})
  prune-local      Keep latest \$KEEP_LOCAL local snapshots on the top dataset (default ${KEEP_LOCAL})
  notify           Send notification email (includes steps, timestamps, backup name & size)

Notes:
  - Run 'init-remote' once before the first upload, or anytime after changing \$HOST_DIR.
  - Local pruning destroys older snapshots **recursively** to maintain consistency across descendants.
  - Set \${BACKUP_EXCLUDE_PROP}=false on a dataset to exclude it (and inherited descendants) from send-to-remote.
    Example: zfs set com.pursotin:backup=false zroot/some/heavy-dataset

Examples:
  $(basename "$0") init-remote snapshot send-to-remote prune-remote prune-local notify
EOF
}

# ----- Main -----
main() {
  if [[ $# -eq 0 ]]; then
    usage
    exit 1
  fi

  # sanity
  [[ -n "${DATASET}" ]] || die "DATASET not set"
  [[ -n "${USER}" && -n "${HOST}" ]] || die "USER/HOST not set"
  [[ -r "${KEYFILE}" ]] || die "KEYFILE not readable: ${KEYFILE}"

  while [[ $# -gt 0 ]]; do
    case "$1" in
      init-remote)    cmd_init_remote ;;
      snapshot)        cmd_snapshot ;;
      send-to-remote)  cmd_send_to_remote ;;
      prune-remote)    cmd_prune_remote ;;
      prune-local)     cmd_prune_local ;;
      notify)          cmd_notify ;;
      -h|--help|help)  usage; exit 0 ;;
      *) die "Unknown subcommand: $1" ;;
    esac
    shift
  done

  echo "Done."
}

main "$@"
{% endraw %}