#!/usr/bin/env bash
set -Eeuo pipefail

###############################################################################
# dockerbackup_prune.sh
#
# Default:
#   A) Create a dated .zst backup locally and keep it
#
# With --apply:
#   A) Create a dated .zst backup
#   C) Prune *.zst under PRUNE_ROOT subfolders, keeping newest N per subfolder
#   D) Keep the new local archive
#
# With --rclone:
#   A) Create a dated .zst backup
#   B) Upload via rclone
#   C) Show what prune would do (unless --apply is also set)
#   D) Keep the new local archive unless --apply is also set
#
# With --rclone --apply:
#   A) Create a dated .zst backup
#   B) Upload via rclone
#   C) Prune for real
#   D) Delete the new local archive
#
# Restore:
#   zstd -d -c /path/backup.zst | tar -x -C /restore/target
###############################################################################

# --------------------------- Defaults (can be overridden) --------------------
SOURCE_DIR="/mybackups"                             # what to archive
PRUNE_ROOT="/mybackups"                             # where to prune .zst
PRUNE_PATTERN="*.zst"                               # files to prune
KEEP_COUNT=5                                        # keep newest N per subfolder

RCLONE_REMOTE="gdrive:location/sublocation"         # rclone destination
OUT_DIR="/tmp"                                      # where to write the .zst
BACKUP_PREFIX="panda-dockerbackup"                  # backup name prefix

# Optional zstd tuning, e.g. "-19 -T0" for max level, multithread
ZSTD_EXTRA_OPTS=""

# --------------------------- Load .env (same dir as script) -----------------
# Order of precedence:
#   defaults < .env (same dir) < CLI flags
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
ENV_FILE="${DOCKERBACKUP_ENV:-$SCRIPT_DIR/.env}"
if [[ -f "$ENV_FILE" ]]; then
  perms="$(stat -c '%a' "$ENV_FILE" 2>/dev/null || stat -f '%Lp' "$ENV_FILE" 2>/dev/null || echo '')"
  if [[ -n "$perms" && ! "$perms" =~ ^(600|400)$ ]]; then
    echo "$(date '+%F %T') [WARN] $ENV_FILE permissions are $perms; recommend 600." >&2
  fi
  set -a
  # shellcheck disable=SC1090
  . "$ENV_FILE"
  set +a
fi

# --------------------------- CLI flags --------------------------------------
APPLY=0
QUIET=0
DRYRUN=0
USE_RCLONE=0

usage() {
  cat <<EOF
Usage: $0 [--rclone] [--apply] [--quiet] [--dry-run] [--keep N] [--root PATH] [--source PATH] [--outdir PATH] [--remote PATH]

Default:
  Creates the local .zst archive and keeps it locally.
  No rclone upload. No prune.

With --apply:
  Creates the local .zst archive, prunes old .zst files for real, and keeps the new archive locally.

With --rclone:
  Uploads the archive via rclone.
  Prune also runs, but only as a dry-run unless --apply is also set.
  The new local archive is kept unless --apply is also set.

With --rclone --apply:
  Uploads the archive, prunes for real, and deletes the new local archive.

Flags:
  --rclone       Enable rclone upload flow
  --apply        Actually delete pruned files; with --rclone also delete the new local archive
  --quiet        Minimal console output (no live progress)
  --dry-run      Simulate everything (no backup, no upload, no deletes)
  --keep N       Keep newest N .zst files per subfolder (default ${KEEP_COUNT})
  --root PATH    Root folder to prune (default ${PRUNE_ROOT})
  --source PATH  Folder to archive (default ${SOURCE_DIR})
  --outdir PATH  Where to write the .zst (default ${OUT_DIR})
  --remote PATH  rclone destination (default ${RCLONE_REMOTE})

Environment file:
  A .env file in the same directory as this script is loaded automatically.
  You can override its path by setting DOCKERBACKUP_ENV=/path/to/.env
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --rclone) USE_RCLONE=1 ;;
    --apply) APPLY=1 ;;
    --quiet) QUIET=1 ;;
    --dry-run) DRYRUN=1 ;;
    --keep) KEEP_COUNT="${2:?}"; shift ;;
    --root) PRUNE_ROOT="${2:?}"; shift ;;
    --source) SOURCE_DIR="${2:?}"; shift ;;
    --outdir) OUT_DIR="${2:?}"; shift ;;
    --remote) RCLONE_REMOTE="${2:?}"; shift ;;
    -h|--help) usage; exit 0 ;;
    *) echo "Unknown option: $1" >&2; usage; exit 2 ;;
  esac
  shift
done

# --------------------------- Logging helpers --------------------------------
ts() { date '+%F %T'; }
log() { [[ $QUIET -eq 0 ]] && echo "$(ts) [$1] $2"; }
log_err() { echo "$(ts) [ERROR] $1" >&2; }
run() { if [[ $DRYRUN -eq 1 ]]; then log "DRY" "$*"; else "$@"; fi; }

truncate_and_print() {
  local s="$1" cols
  if [[ -t 2 ]]; then
    cols=$(tput cols 2>/dev/null || echo 120)
  else
    cols=120
  fi
  (( cols < 20 )) && cols=120
  if (( ${#s} > cols - 1 )); then
    s="${s:0:cols-1}"
  fi
  printf '\r%s\033[K' "$s" >&2
}

cleanup_tmp=""
pipe_pid=""
pipe_pgid=""
pct_fd_open=0
path_fd_open=0
status_line_shown=0

cleanup() {
  local why="${1:-EXIT}"

  if [[ -n "$pipe_pid" ]]; then
    if command -v ps >/dev/null 2>&1; then
      pipe_pgid=$(ps -o pgid= "$pipe_pid" 2>/dev/null | tr -d ' ' || true)
    fi
    if [[ -n "$pipe_pgid" ]]; then
      kill -TERM "-$pipe_pgid" 2>/dev/null || true
      sleep 0.2
      kill -KILL "-$pipe_pgid" 2>/dev/null || true
    else
      kill -TERM "$pipe_pid" 2>/dev/null || true
      sleep 0.2
      kill -KILL "$pipe_pid" 2>/dev/null || true
    fi
    wait "$pipe_pid" 2>/dev/null || true
  fi

  if (( pct_fd_open == 1 )); then exec 4>&- 4<&-; pct_fd_open=0; fi
  if (( path_fd_open == 1 )); then exec 3>&- 3<&-; path_fd_open=0; fi

  if [[ -n "$cleanup_tmp" ]]; then
    rm -f "$cleanup_tmp/pct" "$cleanup_tmp/path" 2>/dev/null || true
    rmdir "$cleanup_tmp" 2>/dev/null || true
  fi

  if (( status_line_shown == 1 )) && [[ $QUIET -eq 0 ]]; then
    printf '\n' >&2
  fi

  case "$why" in
    INT|TERM) exit 130 ;;
  esac
}

trap 'log_err "Command failed: $BASH_COMMAND"; cleanup ERR; exit 1' ERR
trap 'cleanup INT' INT
trap 'cleanup TERM' TERM
trap 'cleanup EXIT' EXIT

# --------------------------- Preconditions ----------------------------------
[[ -d "$SOURCE_DIR" ]] || { log_err "Source dir not found: $SOURCE_DIR"; exit 1; }
[[ -d "$PRUNE_ROOT" ]] || { log_err "Prune root not found: $PRUNE_ROOT"; exit 1; }
command -v tar >/dev/null 2>&1 || { log_err "tar not found"; exit 1; }
command -v zstd >/dev/null 2>&1 || { log_err "zstd not found"; exit 1; }
command -v find >/dev/null 2>&1 || { log_err "find not found"; exit 1; }

if (( USE_RCLONE == 1 )) && (( DRYRUN == 0 )); then
  command -v rclone >/dev/null 2>&1 || { log_err "rclone not found"; exit 1; }
fi

STAMP="${STAMP_OVERRIDE:-$(date +%F_%H-%M)}"
BACKUP_NAME="${BACKUP_PREFIX}-${STAMP}.zst"
BACKUP_PATH="${OUT_DIR%/}/${BACKUP_NAME}"

RCLONE_FLAGS=()
[[ $QUIET -eq 1 ]] && RCLONE_FLAGS+=("-q") || RCLONE_FLAGS+=("--progress")

ZSTD_FLAGS=("-q")
[[ -n "$ZSTD_EXTRA_OPTS" ]] && ZSTD_FLAGS+=($ZSTD_EXTRA_OPTS)

log "INFO" "Mode: $( [[ $DRYRUN -eq 1 ]] && echo 'DRY-RUN ALL' || ([[ $APPLY -eq 1 ]] && echo 'APPLY DELETES' || echo 'NO DELETES') )"
log "INFO" "Rclone mode: $( [[ $USE_RCLONE -eq 1 ]] && echo 'ENABLED' || echo 'DISABLED' )"
log "INFO" "Creating archive: $BACKUP_PATH (.zst)"
log "INFO" "Keeping newest $KEEP_COUNT file(s) per subfolder matching: $PRUNE_PATTERN"
log "INFO" "Source: $SOURCE_DIR"
if (( USE_RCLONE == 1 )); then
  log "INFO" "Remote: $RCLONE_REMOTE"
fi
log "INFO" "Prune root (root itself will be skipped): $PRUNE_ROOT"

# --------------------------- Step A: create .zst with one-liner --------------
log "STEP" "A) Creating archive"

if [[ $DRYRUN -eq 1 ]]; then
  log "DRY" "tar -C \"$SOURCE_DIR\" -cvf - . | pv -n -s \$(du -sb \"$SOURCE_DIR\"|cut -f1) 2>pct.fifo | zstd ${ZSTD_FLAGS[*]} -o \"$BACKUP_PATH\""
else
  mkdir -p "$OUT_DIR"

  export LC_ALL=C
  TOTAL_FILES=$(find "$SOURCE_DIR" -type f | wc -l | awk '{print $1}')
  TOTAL_BYTES=$(du -sb "$SOURCE_DIR" | cut -f1)

  tmpdir="$(mktemp -d)"; cleanup_tmp="$tmpdir"
  path_fifo="$tmpdir/path"; pct_fifo="$tmpdir/pct"
  mkfifo "$path_fifo"
  have_pv=0
  if command -v pv >/dev/null 2>&1 && [[ $QUIET -eq 0 ]] && [[ -t 2 ]]; then
    mkfifo "$pct_fifo"; have_pv=1
  fi

  exec 3<> "$path_fifo"; path_fd_open=1
  if (( have_pv == 1 )); then exec 4<> "$pct_fifo"; pct_fd_open=1; fi

  if (( have_pv == 1 )); then
    (
      set -o pipefail
      tar -C "$SOURCE_DIR" -cvf - . \
        2> >(sed -u 's#^\./##' > "$path_fifo" || true) \
      | pv -n -s "$TOTAL_BYTES" 2> "$pct_fifo" \
      | zstd "${ZSTD_FLAGS[@]}" -o "$BACKUP_PATH"
    ) &
  else
    (
      set -o pipefail
      tar -C "$SOURCE_DIR" -cvf - . \
        2> >(sed -u 's#^\./##' > "$path_fifo" || true) \
      | zstd "${ZSTD_FLAGS[@]}" -o "$BACKUP_PATH"
    ) &
  fi
  pipe_pid=$!

  files_seen=0
  last_path=""
  last_pct="0.0"
  status_line_shown=0
  last_print_time=0

  while kill -0 "$pipe_pid" 2>/dev/null; do
    if (( have_pv == 1 )); then
      if IFS= read -r -t 0.03 pct_line <&4; then last_pct="$pct_line"; fi
    fi

    if IFS= read -r -t 0.03 path_line <&3; then
      last_path="$path_line"
      [[ -n "$last_path" ]] && (( files_seen++ )) || true
    fi

    display_rel="$last_path"
    dir_part="$(dirname -- "$display_rel" 2>/dev/null || echo .)"
    base_part="$(basename -- "$display_rel" 2>/dev/null || echo .)"
    [[ "$dir_part" == "." ]] && dir_part="."

    if [[ $QUIET -eq 0 ]]; then
      now_ns=$(date +%s%N)
      if (( now_ns - last_print_time > 50000000 )); then
        if (( have_pv == 1 )); then
          line="[$(printf '%.1f' "${last_pct:-0}")%] archiving: ${dir_part}/ ${base_part}"
        else
          line="[${files_seen}/${TOTAL_FILES}] archiving: ${dir_part}/ ${base_part}"
        fi
        truncate_and_print "$line"
        status_line_shown=1
        last_print_time=$now_ns
      fi
    fi
  done

  while IFS= read -r -t 0.03 path_line <&3; do
    last_path="$path_line"; (( files_seen++ )) || true
    [[ $QUIET -eq 0 ]] && truncate_and_print "[${files_seen}/${TOTAL_FILES}] archiving: $(dirname -- "$last_path")/ $(basename -- "$last_path")"
    status_line_shown=1
  done || true

  wait "$pipe_pid"
  if [[ $QUIET -eq 0 ]]; then
    printf '\n' >&2
    status_line_shown=0
  fi

  exec 3>&- 3<&-; path_fd_open=0
  if (( have_pv == 1 )); then exec 4>&- 4<&-; pct_fd_open=0; fi
  rm -f "$path_fifo" "$pct_fifo" 2>/dev/null || true
  rmdir "$tmpdir" 2>/dev/null || true
  cleanup_tmp=""
fi
log "OK" "Archive step completed"

# --------------------------- Step B: optional rclone copy --------------------
if (( USE_RCLONE == 1 )); then
  log "STEP" "B) rclone copy to remote"
  if [[ $DRYRUN -eq 1 ]]; then
    log "DRY" "rclone copy \"$BACKUP_PATH\" \"$RCLONE_REMOTE\" ${RCLONE_FLAGS[*]}"
  else
    run rclone copy "$BACKUP_PATH" "$RCLONE_REMOTE" "${RCLONE_FLAGS[@]}"
  fi
  log "OK" "Upload step completed"
else
  log "SKIP" "--rclone not set; skipping upload"
fi

# --------------------------- Step C: prune if apply or rclone ---------------
if (( APPLY == 1 || USE_RCLONE == 1 )); then
  log "STEP" "C) Pruning $PRUNE_PATTERN under $PRUNE_ROOT (keep $KEEP_COUNT)"
  export LC_ALL=C
  prune_ok=1

  while IFS= read -r -d '' dir; do
    mapfile -d '' -t files < <(
      find "$dir" -maxdepth 1 -type f -name "$PRUNE_PATTERN" -printf '%T@ %p\0' |
        sort -z -n |
        cut -z -d' ' -f2-
    )
    count=${#files[@]}
    if (( count == 0 )); then
      log "SKIP" "No matches in: $dir"
      continue
    fi

    keep_from=$(( count - KEEP_COUNT ))
    (( keep_from < 0 )) && keep_from=0

    if (( count <= KEEP_COUNT )); then
      log "KEEP" "At or below limit in: $dir (count=$count)"
      continue
    fi

    keep_list=("${files[@]:keep_from:KEEP_COUNT}")
    del_list=("${files[@]:0:keep_from}")

    for k in "${keep_list[@]}"; do
      log "KEEP" "Keep: $k"
    done
    for f in "${del_list[@]}"; do
      if [[ $DRYRUN -eq 1 || $APPLY -eq 0 ]]; then
        log "WOULD" "Remove: $f"
      else
        if rm -f -- "$f"; then
          log "DELETE" "Removed: $f"
        else
          log_err "Failed to remove: $f"; prune_ok=0
        fi
      fi
    done
  done < <(find "$PRUNE_ROOT" -mindepth 1 -type d -print0)

  if [[ $prune_ok -ne 1 ]]; then
    log_err "Prune step encountered errors; aborting final delete"
    exit 1
  fi
  log "OK" "Prune step completed"
else
  log "SKIP" "--apply and --rclone not set; skipping prune"
fi

# --------------------------- Step D: remove local .zst -----------------------
if (( USE_RCLONE == 1 )); then
  log "STEP" "D) Removing local backup"
  if [[ $DRYRUN -eq 1 || $APPLY -eq 0 ]]; then
    log "WOULD" "Remove local backup: $BACKUP_PATH"
  else
    if rm -f -- "$BACKUP_PATH"; then
      log "DELETE" "Removed local backup: $BACKUP_PATH"
    else
      log_err "Failed to remove local backup: $BACKUP_PATH"
      exit 1
    fi
  fi
else
  log "KEEP" "Local archive retained: $BACKUP_PATH"
fi

log "DONE" "All steps completed successfully"