compress.sh

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

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]:-$0}")" && pwd)"
# shellcheck source=common.sh
source "${SCRIPT_DIR}/common.sh"

# Defaults
SCAN_DIR="."
SHA1_FILE=""
SHA1_APPEND=0
SHA256_FILE=""
SHA256_APPEND=0
CHECKSUM_DELIM=$'\x1f'

THRESHOLD_BYTES=$((100 * 1024 * 1024)) # 100 MiB
SMALL_JOBS=8
BIG_JOBS=1

SMALL_COMPRESSOR="xz" # xz | zstd
BIG_COMPRESSOR="pixz" # pixz | xz | pzstd | zstd

XZ_LEVEL="$DEFAULT_XZ_LEVEL"
ZSTD_LEVEL="$DEFAULT_ZSTD_LEVEL"
PZSTD_LEVEL="$DEFAULT_PZSTD_LEVEL"

QUIET=0
KEEP_SOURCE=0
EXTENSIONS=(tar sql txt csv ibd xlsx docx log)
EXTENSIONS_CUSTOMIZED=0

usage() {
  cat >&2 <<'EOF'
Usage:
  compress.sh [options]

Options:
  -d, --dir DIR           Directory to scan (default: .)
  -s, --sha1 FILE         Create/truncate FILE and write SHA1 of originals (before compression per file)
      --sha1-append       Append to FILE instead of truncating
      --sha256 FILE       Create/truncate FILE and write SHA256 of originals (before compression per file)
      --sha256-append     Append to FILE instead of truncating
  -t, --threshold SIZE    Small/Big split (default: 100MiB). Examples: 100M, 200MiB, 50000000
  -j, --jobs N            Parallel jobs for small files (default: 8)

      --small TOOL        Small-file compressor: xz or zstd (default: xz)
      --big TOOL          Big-file compressor: pixz, xz, pzstd, or zstd (default: pixz)
      --big-jobs N        Parallel jobs for big files (default: 1)

      --xz-level  -#      xz/pixz level (default: -5)
      --zstd-level -#     zstd level (default: -6)
      --pzstd-level -#    pzstd level (default: -10)

      --ext EXT           Add an extension (without leading dot) to the scan list.
                          May be repeated; first use replaces the defaults.

  -k, --keep              Keep original files after successful compression
                          (default: originals are removed once the compressed
                          output is in place)
  -q, --quiet             Less logging
  -h, --help              Show help
EOF
}

parse_size() {
  local s="${1,,}" n unit
  if [[ "$s" =~ ^[0-9]+$ ]]; then
    echo "$s"
    return 0
  fi
  if [[ "$s" =~ ^([0-9]+)(k|kb|kib)$ ]]; then
    n="${BASH_REMATCH[1]}"
    echo $((n * 1024))
    return 0
  fi
  if [[ "$s" =~ ^([0-9]+)(m|mb|mib)$ ]]; then
    n="${BASH_REMATCH[1]}"
    echo $((n * 1024 * 1024))
    return 0
  fi
  if [[ "$s" =~ ^([0-9]+)(g|gb|gib)$ ]]; then
    n="${BASH_REMATCH[1]}"
    echo $((n * 1024 * 1024 * 1024))
    return 0
  fi
  echo "Invalid size: $1" >&2
  exit 2
}

emit_checksum_lines() {
  local sha1_line="${1-}" sha256_line="${2-}"
  if [[ -n "$sha1_line" ]]; then
    printf 'sha1%s%s\n' "$CHECKSUM_DELIM" "$sha1_line"
  fi
  if [[ -n "$sha256_line" ]]; then
    printf 'sha256%s%s\n' "$CHECKSUM_DELIM" "$sha256_line"
  fi
}

append_checksum_line() {
  local algo="$1" line="$2"
  case "$algo" in
  sha1)
    [[ -n "$SHA1_FILE" ]] && printf '%s\n' "$line" >>"$SHA1_FILE"
    ;;
  sha256)
    [[ -n "$SHA256_FILE" ]] && printf '%s\n' "$line" >>"$SHA256_FILE"
    ;;
  esac
}

fanout_checksums() {
  local line algo payload
  while IFS= read -r line; do
    [[ -z "$line" ]] && continue
    IFS=$CHECKSUM_DELIM read -r algo payload <<<"$line"
    append_checksum_line "$algo" "$payload"
  done
}

# --- arg parsing ---
while [[ $# -gt 0 ]]; do
  case "$1" in
  -d | --dir)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    SCAN_DIR="$2"
    shift 2
    ;;
  -d?*)
    SCAN_DIR="${1:2}"
    shift
    ;;
  -s | --sha1)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    SHA1_FILE="$2"
    shift 2
    ;;
  -s?*)
    SHA1_FILE="${1:2}"
    shift
    ;;
  --sha1-append)
    SHA1_APPEND=1
    shift
    ;;
  --sha256)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    SHA256_FILE="$2"
    shift 2
    ;;
  --sha256-append)
    SHA256_APPEND=1
    shift
    ;;
  -t | --threshold)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    THRESHOLD_BYTES="$(parse_size "$2")"
    shift 2
    ;;
  -t?*)
    THRESHOLD_BYTES="$(parse_size "${1:2}")"
    shift
    ;;
  -j | --jobs)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    SMALL_JOBS="$2"
    shift 2
    ;;
  -j?*)
    SMALL_JOBS="${1:2}"
    shift
    ;;
  --small)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    SMALL_COMPRESSOR="$2"
    shift 2
    ;;
  --big)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    BIG_COMPRESSOR="$2"
    shift 2
    ;;
  --big-jobs)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    BIG_JOBS="$2"
    shift 2
    ;;
  --xz-level)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    XZ_LEVEL="$2"
    shift 2
    ;;
  --zstd-level)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    ZSTD_LEVEL="$2"
    shift 2
    ;;
  --pzstd-level)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    PZSTD_LEVEL="$2"
    shift 2
    ;;
  --ext)
    [[ $# -lt 2 ]] && die "Missing value for $1"
    if [[ "$EXTENSIONS_CUSTOMIZED" -eq 0 ]]; then
      EXTENSIONS=()
      EXTENSIONS_CUSTOMIZED=1
    fi
    IFS=',' read -r -a _ext_parts <<<"${2// /,}"
    for part in "${_ext_parts[@]}"; do
      part="${part,,}"
      part="${part#.}"
      [[ -z "$part" ]] && continue
      EXTENSIONS+=("$part")
    done
    shift 2
    ;;
  -k | --keep | --keep-source)
    KEEP_SOURCE=1
    shift
    ;;
  -q | --quiet)
    QUIET=1
    shift
    ;;
  -h | --help)
    usage
    exit 0
    ;;
  --)
    shift
    break
    ;;
  *)
    echo "Unknown option: $1" >&2
    usage
    exit 2
    ;;
  esac
done

[[ $# -eq 0 ]] || die "Unexpected positional arguments: $*"
[[ -d "$SCAN_DIR" ]] || die "Scan directory not found: $SCAN_DIR"
[[ "$SMALL_JOBS" =~ ^[0-9]+$ && "$SMALL_JOBS" -ge 1 ]] || die "--jobs must be a positive integer (got: $SMALL_JOBS)"
[[ "$BIG_JOBS" =~ ^[0-9]+$ && "$BIG_JOBS" -ge 1 ]] || die "--big-jobs must be a positive integer (got: $BIG_JOBS)"

# Validate required tools before doing any work
require_cmd parallel
require_cmd file
require_cmd pv
case "$SMALL_COMPRESSOR" in
xz) require_cmd xz ;;
zstd) require_cmd zstd ;;
*) die "SMALL_COMPRESSOR must be xz or zstd" ;;
esac
case "$BIG_COMPRESSOR" in
pixz) require_cmd pixz ;;
xz) require_cmd xz ;;
pzstd) require_cmd pzstd ;;
zstd) require_cmd zstd ;;
*) die "BIG_COMPRESSOR must be pixz, xz, pzstd, or zstd" ;;
esac
[[ -n "$SHA1_FILE" ]] && require_cmd sha1sum
[[ -n "$SHA256_FILE" ]] && require_cmd sha256sum

# Create/truncate checksum files up front (if requested)
if [[ -n "$SHA256_FILE" && "$SHA256_APPEND" -eq 1 && -e "$SHA256_FILE" && ! -s "$SHA256_FILE" ]]; then
  log "Detected empty SHA-256 manifest; overwriting $SHA256_FILE"
  SHA256_APPEND=0
fi
prepare_file_for_write "$SHA1_FILE" "$SHA1_APPEND"
prepare_file_for_write "$SHA256_FILE" "$SHA256_APPEND"

skip_if_already_compressed() {
  local file="$1" actual
  actual="$(detect_actual_format "$file")"
  case "$actual" in
  gz | bz2 | xz | zst)
    log "skip (already compressed: ${actual}): $file"
    return 0
    ;;
  esac
  return 1
}

out_name() {
  local f="$1" algo="$2"
  local ext
  ext="$(get_compression_extension "$algo")"
  if [[ "$ext" == "unknown" ]]; then
    echo "Unknown compressor: $algo" >&2
    return 2
  fi
  printf '%s\n' "${f}.${ext}"
}

compress_small() {
  # Job shells spawned by parallel run `bash -c` WITHOUT errexit/pipefail:
  # re-enable them here so a failed reader/compressor aborts the pipeline and
  # a failed mv does not fall through to deleting the source file.
  set -euo pipefail
  local f="$1" out tmp sha1_line="" sha256_line="" size=""
  if skip_if_already_compressed "$f"; then
    return 0
  fi
  out="$(out_name "$f" "$SMALL_COMPRESSOR")"
  [[ -e "$out" ]] && {
    log "skip: $f"
    return 0
  }

  tmp="${out}.tmp.${PARALLEL_SEQ:-$$}"
  rm -f -- "$tmp"

  local sha1_tmp="" sha256_tmp=""
  [[ -n "${SHA1_FILE:-}" ]] && sha1_tmp="$(mktemp)"
  if [[ -n "${SHA256_FILE:-}" ]]; then
    sha256_tmp="$(mktemp)"
    size="$(stat -c '%s' -- "$f" 2>/dev/null || stat -f '%z' -- "$f")"
  fi

  local -a read_cmd=(cat -- "$f")
  local compress_cmd_str

  log "small(${SMALL_COMPRESSOR}): $f -> $out"

  compress_cmd_str=$(get_compressor_command "$SMALL_COMPRESSOR" compress)
  if [[ "$compress_cmd_str" == "unknown" ]]; then
    echo "SMALL_COMPRESSOR must be xz or zstd" >&2
    rm -f -- "$tmp"
    [[ -n "$sha1_tmp" ]] && rm -f -- "$sha1_tmp"
    [[ -n "$sha256_tmp" ]] && rm -f -- "$sha256_tmp"
    return 2
  fi

  # Hash taps: tee >(cmd) process substitutions are untracked children in
  # bash (no $!, bare wait does not reap them), so each checksum job reads a
  # fifo as a real background job that can be waited on explicitly.
  local -a tap_fifos=() tap_pids=()
  local fifo
  if [[ -n "$sha1_tmp" ]]; then
    fifo="$(mktemp -u)"
    mkfifo -- "$fifo"
    sha1sum <"$fifo" | awk -v path="$f" '{printf "%s  %s\n", $1, path}' >"$sha1_tmp" &
    tap_pids+=("$!")
    tap_fifos+=("$fifo")
  fi
  if [[ -n "$sha256_tmp" ]]; then
    fifo="$(mktemp -u)"
    mkfifo -- "$fifo"
    sha256sum <"$fifo" | awk -v path="$f" -v size="$size" '{printf "%s  %s  %s\n", $1, path, size}' >"$sha256_tmp" &
    tap_pids+=("$!")
    tap_fifos+=("$fifo")
  fi

  local pipeline_failed=0
  if ((${#tap_fifos[@]})); then
    if ! "${read_cmd[@]}" | tee -- "${tap_fifos[@]}" | $compress_cmd_str >"$tmp"; then
      pipeline_failed=1
    fi
  else
    if ! "${read_cmd[@]}" | $compress_cmd_str >"$tmp"; then
      pipeline_failed=1
    fi
  fi

  # Reap the hash taps before reading their output files.
  local tap_rc=0 pid
  for pid in "${tap_pids[@]}"; do
    wait "$pid" || tap_rc=1
  done
  if ((${#tap_fifos[@]})); then
    rm -f -- "${tap_fifos[@]}"
  fi

  if ((pipeline_failed || tap_rc)); then
    rm -f -- "$tmp"
    [[ -n "$sha1_tmp" ]] && rm -f -- "$sha1_tmp"
    [[ -n "$sha256_tmp" ]] && rm -f -- "$sha256_tmp"
    return 1
  fi

  if [[ -n "$sha1_tmp" ]]; then
    sha1_line="$(cat "$sha1_tmp")"
    rm -f -- "$sha1_tmp"
  fi
  if [[ -n "$sha256_tmp" ]]; then
    sha256_line="$(cat "$sha256_tmp")"
    rm -f -- "$sha256_tmp"
  fi

  # A requested checksum that produced nothing is a job failure — emitting a
  # manifest line without a hash would be worse than failing loudly.
  if [[ -n "${SHA1_FILE:-}" && -z "$sha1_line" ]] || [[ -n "${SHA256_FILE:-}" && -z "$sha256_line" ]]; then
    rm -f -- "$tmp"
    return 1
  fi

  touch -r "$f" "$tmp"
  mv -f -- "$tmp" "$out"
  [[ "$KEEP_SOURCE" -eq 0 ]] && rm -f -- "$f"

  emit_checksum_lines "$sha1_line" "$sha256_line"
  return 0
}

compress_big_seq() {
  set -euo pipefail
  local f="$1" out tmp sha1_line="" sha256_line="" size=""
  if skip_if_already_compressed "$f"; then
    return 0
  fi
  out="$(out_name "$f" "$BIG_COMPRESSOR")"
  [[ -e "$out" ]] && {
    log "skip: $f"
    return 0
  }

  tmp="${out}.tmp.$$"
  rm -f -- "$tmp"

  local sha1_tmp="" sha256_tmp=""
  [[ -n "$SHA1_FILE" ]] && sha1_tmp="$(mktemp)"
  if [[ -n "$SHA256_FILE" ]]; then
    sha256_tmp="$(mktemp)"
    size="$(stat -c '%s' -- "$f" 2>/dev/null || stat -f '%z' -- "$f")"
  fi

  local -a read_cmd=(pv -ptebar -N "$f" -- "$f")

  log "big(${BIG_COMPRESSOR}): $f -> $out"

  # xz/zstd selected for the big path use their multithreaded variants
  # (the plain arms are single-threaded on purpose, for small-file jobs).
  local big_comp="$BIG_COMPRESSOR"
  [[ "$big_comp" == "xz" ]] && big_comp="xz_big"
  [[ "$big_comp" == "zstd" ]] && big_comp="zstd_big"

  local compress_cmd_str
  compress_cmd_str=$(get_compressor_command "$big_comp" compress)
  if [[ "$compress_cmd_str" == "unknown" ]]; then
    echo "BIG_COMPRESSOR must be pixz, xz, pzstd, or zstd" >&2
    rm -f -- "$tmp"
    [[ -n "$sha1_tmp" ]] && rm -f -- "$sha1_tmp"
    [[ -n "$sha256_tmp" ]] && rm -f -- "$sha256_tmp"
    return 2
  fi

  # Hash taps via fifo + real background job — process substitutions from
  # tee >(cmd) cannot be waited on (see compress_small).
  local -a tap_fifos=() tap_pids=()
  local fifo
  if [[ -n "$sha1_tmp" ]]; then
    fifo="$(mktemp -u)"
    mkfifo -- "$fifo"
    sha1sum <"$fifo" | awk -v path="$f" '{printf "%s  %s\n", $1, path}' >"$sha1_tmp" &
    tap_pids+=("$!")
    tap_fifos+=("$fifo")
  fi
  if [[ -n "$sha256_tmp" ]]; then
    fifo="$(mktemp -u)"
    mkfifo -- "$fifo"
    sha256sum <"$fifo" | awk -v path="$f" -v size="$size" '{printf "%s  %s  %s\n", $1, path, size}' >"$sha256_tmp" &
    tap_pids+=("$!")
    tap_fifos+=("$fifo")
  fi

  local pipeline_failed=0
  if ((${#tap_fifos[@]})); then
    if ! "${read_cmd[@]}" | tee -- "${tap_fifos[@]}" | $compress_cmd_str >"$tmp"; then
      pipeline_failed=1
    fi
  else
    if ! "${read_cmd[@]}" | $compress_cmd_str >"$tmp"; then
      pipeline_failed=1
    fi
  fi

  local tap_rc=0 pid
  for pid in "${tap_pids[@]}"; do
    wait "$pid" || tap_rc=1
  done
  if ((${#tap_fifos[@]})); then
    rm -f -- "${tap_fifos[@]}"
  fi

  if ((pipeline_failed || tap_rc)); then
    rm -f -- "$tmp"
    [[ -n "$sha1_tmp" ]] && rm -f -- "$sha1_tmp"
    [[ -n "$sha256_tmp" ]] && rm -f -- "$sha256_tmp"
    return 1
  fi

  if [[ -n "$sha1_tmp" ]]; then
    sha1_line="$(cat "$sha1_tmp")"
    rm -f -- "$sha1_tmp"
  fi
  if [[ -n "$sha256_tmp" ]]; then
    sha256_line="$(cat "$sha256_tmp")"
    rm -f -- "$sha256_tmp"
  fi

  if [[ -n "${SHA1_FILE:-}" && -z "$sha1_line" ]] || [[ -n "${SHA256_FILE:-}" && -z "$sha256_line" ]]; then
    rm -f -- "$tmp"
    return 1
  fi

  touch -r "$f" "$tmp"
  mv -f -- "$tmp" "$out"
  [[ "$KEEP_SOURCE" -eq 0 ]] && rm -f -- "$f"

  [[ -n "$sha1_line" ]] && append_checksum_line sha1 "$sha1_line"
  [[ -n "$sha256_line" ]] && append_checksum_line sha256 "$sha256_line"
  return 0
}

export -f compress_small compress_big_seq out_name emit_checksum_lines append_checksum_line log skip_if_already_compressed detect_actual_format get_expected_extension get_compression_extension get_compressor_command
export SMALL_COMPRESSOR BIG_COMPRESSOR BIG_JOBS XZ_LEVEL ZSTD_LEVEL PZSTD_LEVEL QUIET SHA1_FILE SHA256_FILE CHECKSUM_DELIM KEEP_SOURCE

small_list="$(mktemp)"
big_list="$(mktemp)"
trap 'rm -f "$small_list" "$big_list"' EXIT

# collect + split
if [[ "${#EXTENSIONS[@]}" -eq 0 ]]; then
  echo "No extensions configured; nothing to do." >&2
  exit 0
fi

ext_pred=()
for i in "${!EXTENSIONS[@]}"; do
  ext="${EXTENSIONS[$i]}"
  [[ -z "$ext" ]] && continue
  if [[ "$i" -gt 0 ]]; then
    ext_pred+=("-o")
  fi
  ext_pred+=("-iname" "*.${ext}")
done

if [[ "${#ext_pred[@]}" -eq 0 ]]; then
  echo "No valid extensions configured; nothing to do." >&2
  exit 0
fi

find "$SCAN_DIR" -type f \( "${ext_pred[@]}" \) \
  ! -iname '*.zst' ! -iname '*.tzst' ! -iname '*.xz' ! -iname '*.txz' \
  ! -iname '*.gz' ! -iname '*.tgz' ! -iname '*.bz2' ! -iname '*.tbz' ! -iname '*.tbz2' -print0 |
  LC_ALL=C sort -z |
  while IFS= read -r -d '' f; do
    # Never process the checksum outputs themselves (-ef also covers
    # path variants like ./x vs x and hardlinks to the same file).
    if [[ -n "$SHA1_FILE" && -e "$SHA1_FILE" && "$f" -ef "$SHA1_FILE" ]]; then
      continue
    fi
    if [[ -n "$SHA256_FILE" && -e "$SHA256_FILE" && "$f" -ef "$SHA256_FILE" ]]; then
      continue
    fi
    if [[ "$(stat -c '%s' -- "$f" 2>/dev/null || stat -f '%z' -- "$f")" -lt "$THRESHOLD_BYTES" ]]; then
      printf '%s\0' "$f" >>"$small_list"
    else
      printf '%s\0' "$f" >>"$big_list"
    fi
  done || true

# GNU parallel's --bar writes to /dev/tty and errors in non-interactive
# environments; only enable it when a usable controlling tty exists.
parallel_bar=()
if [[ -t 1 ]] && (: </dev/tty) 2>/dev/null; then
  parallel_bar=(--bar)
fi

# small files in parallel
if [[ -n "$SHA1_FILE" || -n "$SHA256_FILE" ]]; then
  parallel -0 --no-run-if-empty "${parallel_bar[@]}" --lb -j "$SMALL_JOBS" compress_small {} :::: "$small_list" | fanout_checksums
else
  parallel -0 --no-run-if-empty "${parallel_bar[@]}" -j "$SMALL_JOBS" compress_small {} :::: "$small_list"
fi

# big files (parallelizable)
if [[ -s "$big_list" ]]; then
  if [[ -n "$SHA1_FILE" || -n "$SHA256_FILE" ]]; then
    parallel -0 --no-run-if-empty "${parallel_bar[@]}" --lb -j "$BIG_JOBS" compress_big_seq {} :::: "$big_list"
  else
    parallel -0 --no-run-if-empty "${parallel_bar[@]}" -j "$BIG_JOBS" compress_big_seq {} :::: "$big_list"
  fi
fi

exit 0