# Sourced by bench/bench (after lib.sh): what the work directory and the store hold (`bench du`), what `bench clean` removes, the peak
# of a run, and the pruning of the caches after a session.
#
# ROOT (BENCH_WORK) holds the live clusters, the scratch and the engine builds; STORE (BENCH_STORE, default ROOT) the datasets, the
# downloads, the corpus dumps and the cache of lanes. Everything a run creates is either a cache (regenerated by the next run) or
# scratch (removed when the run ends).

du_kb() { du -sk "$@" 2>/dev/null | awk '{s += $1} END {print s + 0}'; }
du_roots() { if [[ $STORE == "$ROOT" ]]; then du_kb "$ROOT"; else du_kb "$ROOT" "$STORE"; fi; }   # the whole footprint
human() { awk -v k="$1" 'BEGIN {if (k >= 1048576) printf "%.2f GiB", k / 1048576; else if (k >= 1024) printf "%.0f MiB", k / 1024; else printf "%d KiB", k}'; }
fs_of() { stat -f -c %T "$1" 2>/dev/null || echo "?"; }

# footprint_rows: `KiB<TAB>component<TAB>what` for everything under ROOT and STORE, in the order of the tables of the README
footprint_rows() {
  local d n
  row() { [[ -e $2 ]] && printf '%s\t%s\t%s\n' "$(du_kb "$2")" "$1" "$3" || true; }
  for d in "$DATA"/articles.tsv*; do row "data/$(basename "$d")" "$d" 'Simple English Wikipedia, converted (setup)'; done
  for d in "$DATA"/places-*; do row "data/$(basename "$d")" "$d" 'GeoNames places and queries (setup)'; done
  for d in "$DATA"/Q "$DATA"/F; do row "data/$(basename "$d")" "$d" 'corpus exported for the engines outside Postgres (derived)'; done
  row 'download/' "$DOWNLOADS" 'raw downloads (removed after setup unless KEEP_DOWNLOADS=1)'
  for d in "$ENG"/*/; do
    [[ -d $d ]] || continue
    n=$(basename "$d")
    printf '%s\t%s\t%s\n' "$(du -sk --exclude=Q --exclude=F "$d" 2>/dev/null | cut -f1)" "engines/$n" 'built or unpacked engine'
    for t in Q F; do row "engines/$n/$t" "$d$t" 'data of an engine outside Postgres, removed with the run'; done
  done
  row 'target/' "$ROOT/target" 'xbench binary'
  row 'build/' "$ROOT/build" 'build trees (removed after setup)'
  printf '%s\t%s\t%s\n' "$(du_kb "$CACHE"/*.dump)" 'cache/*.dump' 'corpus dumps, one per tier'
  row 'cache/lanes' "$CACHE/lanes" "$(lane_count) cached lanes (compressed, kept with --keep)"
  row 'pg/' "$ROOT/pg" 'live clusters of the run'
  for d in run articles places; do row "$d/" "$ROOT/$d" 'scratch of the run'; done
}
lane_count() { find "$CACHE/lanes" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | wc -l; }

cmd_du() {
  echo "work   $ROOT ($(fs_of "$ROOT"), $(df -h --output=avail "$ROOT" 2>/dev/null | tail -1 | tr -d ' ') free)"
  [[ $STORE == "$ROOT" ]] || echo "store  $STORE ($(fs_of "$STORE"), $(df -h --output=avail "$STORE" 2>/dev/null | tail -1 | tr -d ' ') free)"
  footprint_rows | awk -F'\t' '{printf "%12s  %-26s %s\n", sprintf("%d", $1), $2, $3}' | while read -r kb rest; do printf '%12s  %s\n' "$(human "$kb")" "$rest"; done
  printf '%12s  total\n' "$(human "$(du_roots)")"
}

# --- clean: the scratch always; --cache the dumps, the cache of lanes and the derived exports; --build the engines; --all everything
cmd_clean() {
  local t
  for t in Q F; do TIER=$t "$STEPS/cluster.sh" destroy 2>/dev/null || true; done
  rm -rf "$ROOT/pg" "$ROOT/articles" "$ROOT/places" "$ROOT/run" "$ROOT/build"
  for t in "$ENG"/*/Q "$ENG"/*/F; do rm -rf "$t"; done   # the data of an engine outside Postgres
  [[ -z $clean_cache && -z $all_clean ]] || rm -rf "$CACHE" "$DATA/Q" "$DATA/F"
  [[ -z $clean_build && -z $all_clean ]] || rm -rf "$ENG" "$ROOT/target"
  [[ -z $all_clean ]] || rm -rf "$DATA" "$DOWNLOADS"
  echo "removed the scratch under $ROOT${clean_cache:+, the dumps and the cache of lanes}${clean_build:+, the engine builds}${all_clean:+, the dumps, the cache of lanes, the engine builds, the datasets and the downloads}"
}

# --- the peak of a run: a sampler of the sizes of ROOT and STORE (a `du` of a few thousand files every 5 s)
peak_start() {
  PEAK_FILE=$(mktemp "${TMPDIR:-/tmp}/bench-peak.XXXXXX")
  echo 0 >"$PEAK_FILE"
  ( renice -n 19 -p "$BASHPID" >/dev/null; while :; do
      now=$(du_roots); [[ $now -le $(<"$PEAK_FILE") ]] || echo "$now" >"$PEAK_FILE"
      sleep 5
    done ) >/dev/null 2>&1 &
  PEAK_PID=$!
}
peak_report() { # the peak and what is left
  [[ -n ${PEAK_PID:-} ]] || return 0
  kill "$PEAK_PID" 2>/dev/null || true; wait "$PEAK_PID" 2>/dev/null || true
  local peak left; peak=$(<"$PEAK_FILE"); left=$(du_roots); (( left <= peak )) || peak=$left
  echo "footprint: peak $(human "$peak") under $ROOT$([[ $STORE == "$ROOT" ]] || echo " and $STORE"), $(human "$left") left (bench du)"
  rm -f "$PEAK_FILE"; PEAK_PID=''
}

lane_gc() { # the cache of lanes at most BENCH_CACHE_MAX MiB: what is not a compressed lane (older layouts, interrupted saves) goes first, then the least recently used
  local lanes=$CACHE/lanes max=$(( ${BENCH_CACHE_MAX:-4096} * 1024 )) d
  [[ -d $lanes ]] || return 0
  for d in "$lanes"/*; do [[ -s $d/data.tar.zst ]] || rm -rf "$d"; done
  while (( $(du_kb "$lanes") > max )); do
    d=$(find "$lanes" -mindepth 1 -maxdepth 1 -printf '%T@ %p\n' | sort -n | head -1 | cut -d' ' -f2-)
    [[ -n $d ]] || break
    rm -rf "$d"
  done
}

# --- after a successful session: the cache of lanes is removed unless it is kept (--keep), else the oldest entries above the limit go
cache_prune() {
  if [[ -z ${KEEP_CACHE:-} ]]; then rm -rf "$CACHE/lanes" "$DATA/Q" "$DATA/F"; return 0; fi
  lane_gc
}
