#!/usr/bin/env bash
# shellcheck disable=SC2154  # variables come from the environment
# hostrepo-heal: detect and repair damage (bit rot, truncation, corruption) in
# a restic repository, with par2 parity and refetching from another copy.
#
# Every restic file except `config` is named after the SHA-256 of its own
# contents (data/, index/, snapshots/, keys/). Damage is detected by hashing a
# file against its name: no parity or network needed. A repair is accepted
# only when the repaired file hashes to its name again, so it is always the
# exact original.
#
#   hostrepo-heal local <repo-dir> <par2-dir> [options]
#       A repository on local disk (a peer mirror). Hash every file; repair a
#       damaged one from its par2 set, else refetch it from --remote; move the
#       damaged copy to --quarantine. With --create, make par2 sets for files
#       that lack one. With --par2-verify K/N, also par2-verify the sets of
#       every Nth file (starting at K), so rot in the parity itself is found;
#       a damaged set is rebuilt from the verified file.
#         --remote RCLONE:PATH   refetch source (the same repository in B2)
#         --quarantine DIR       where damaged copies go (default: <par2-dir>/../quarantine)
#
#   hostrepo-heal b2-sync <HOST>
#       Keep par2 sets for this host's repository in the bucket, under
#       par2/<HOST>/. Downloads each repository file that has no set yet,
#       checks its hash (so damage in B2 is found on the day it is written),
#       creates and uploads the set. Deletes sets whose file restic pruned.
#
#   hostrepo-heal b2-lite <HOST> [--repair]
#       The cheap daily check: list the repository with the SHA-1 B2 stores for
#       each file (no download) and compare with SHA1SUMS, the SHA-1s recorded
#       when each file was proven good. With --repair, repair any mismatch.
#
#   hostrepo-heal b2-repair <HOST> <id>...
#       Repair the named repository files in B2 (as named by a failed
#       `restic check`). Downloads each file and its set, repairs it if its
#       hash is wrong, and uploads the exact original. Ids that are not files
#       in the repository are ignored.
#
# B2 access for b2-*: B2_ACCOUNT_ID, B2_ACCOUNT_KEY and RESTIC_REPOSITORY
# (b2:<bucket>:<HOST>) from the environment, as in /etc/hostrepo/hostrepo.env.
# Common options: --redundancy PCT (default 10), --dry-run.
# Exit: 0 clean, 1 damage found and repaired (or sets rebuilt), 2 not repaired.
set -euo pipefail
export LC_ALL=C
PAR2="${PAR2_BIN:-par2}"
RCLONE="${RCLONE_BIN:-$(command -v rclone || true)}"
REDUNDANCY="${HOSTREPO_PAR2_REDUNDANCY:-10}"
STATUS=0; DRY=0
log()  { printf '%s %s\n' "$(date -u +%FT%TZ)" "$*" >&2; }
warn() { [[ $STATUS -lt 1 ]] && STATUS=1; log "REPAIRED: $*"; }
fail() { STATUS=2; log "NOT REPAIRED: $*"; }
die()  { log "FATAL: $*"; exit 2; }
command -v "$PAR2" >/dev/null || die "par2 not installed"

# A path relative to a repository root that is named after its own hash.
addressed() { [[ "$1" =~ ^(data/[0-9a-f]{2}/|index/|snapshots/|keys/)[0-9a-f]{64}$ ]]; }
hash_ok()   { [[ "$(sha256sum < "$1" | cut -c1-64)" == "${2##*/}" ]]; }

# par2_create <root> <rel> <par2-root>: set at <par2-root>/<rel>.par2 (+ one
# volume), and <rel>.sums with the SHA-256 of each set file. par2 verify checks
# the data against the set but says nothing when the set itself has rotted;
# the sums catch that.
par2_create() {
  local d; d="$(dirname "$3/$2")"; mkdir -p "$d"
  rm -f "$3/$2.par2" "$3/$2".vol*.par2 "$3/$2.sums"
  "$PAR2" create -q -q -r"$REDUNDANCY" -n1 -B "$1" "$3/$2.par2" "$1/$2" >/dev/null &&
    (cd "$d" && sha256sum -- "${2##*/}".*par2 > "${2##*/}.sums")
}
# The set's own files are intact, and the data verifies against it.
set_ok() { (cd "$(dirname "$3/$2")" && sha256sum --quiet --status -c "${2##*/}.sums" 2>/dev/null) && par2_verify "$@"; }
par2_verify() { "$PAR2" verify -q -q -B "$1" "$3/$2.par2" >/dev/null 2>&1; }
par2_repair() { "$PAR2" repair -q -q -B "$1" "$3/$2.par2" >/dev/null 2>&1; }

# ---------------------------------------------------------------- local ----
cmd_local() {
  local repo="${1:?repo dir}" pdir="${2:?par2 dir}"; shift 2
  local create=0 k=0 n=0 remote="" quar=""
  while [[ $# -gt 0 ]]; do case "$1" in
    --create) create=1 ;; --par2-verify) k="${2%/*}"; n="${2#*/}"; shift ;;
    --remote) remote="$2"; shift ;; --quarantine) quar="$2"; shift ;;
    --redundancy) REDUNDANCY="$2"; shift ;; --dry-run) DRY=1 ;;
    *) die "unknown option $1" ;; esac; shift; done
  [[ -d "$repo" ]] || die "no repository at $repo"
  quar="${quar:-$(dirname "$pdir")/quarantine}"
  local ts; ts=$(date -u +%Y%m%dT%H%M%SZ)
  local i=0 ok=0 made=0 rel f
  while IFS= read -r rel; do
    f="$repo/$rel"; i=$((i+1))
    addressed "$rel" || continue   # config is checked by restic itself
    if hash_ok "$f" "$rel"; then
      ok=$((ok+1))
      if [[ ! -f "$pdir/$rel.par2" ]]; then
        if [[ $create -eq 1 && $DRY -eq 0 ]]; then par2_create "$repo" "$rel" "$pdir" && made=$((made+1)) || fail "$rel: par2 create failed"; fi
      elif [[ $n -gt 0 ]] && (( (i - k) % n == 0 )) && ! set_ok "$repo" "$rel" "$pdir"; then
        [[ $DRY -eq 1 ]] && { log "dry-run: would rebuild the par2 set of $rel"; continue; }
        par2_create "$repo" "$rel" "$pdir" && warn "$rel: par2 set was damaged; rebuilt from the verified file" || fail "$rel: par2 set damaged and rebuild failed"
      fi
      continue
    fi
    log "DAMAGED: $rel does not match its hash"
    [[ $DRY -eq 1 ]] && { log "dry-run: would repair $rel"; STATUS=2; continue; }
    mkdir -p "$(dirname "$quar/$ts/$rel")"; cp -p "$f" "$quar/$ts/$rel"
    if [[ -f "$pdir/$rel.par2" ]] && par2_repair "$repo" "$rel" "$pdir" && hash_ok "$f" "$rel"; then
      rm -f "$f.1"; warn "$rel: repaired from par2 (damaged copy in $quar/$ts)"; continue
    fi
    [[ -f "$f.1" ]] && mv -f "$f.1" "$f"   # a failed repair may have renamed it
    if [[ -n "$remote" && -n "$RCLONE" ]] && "$RCLONE" copyto -q "$remote/$rel" "$f.heal" && hash_ok "$f.heal" "$rel"; then
      mv -f "$f.heal" "$f"; par2_create "$repo" "$rel" "$pdir" || true
      warn "$rel: refetched from $remote (damaged copy in $quar/$ts)"; continue
    fi
    rm -f "$f.heal"; fail "$rel: par2 and refetch both failed; damaged copy in $quar/$ts"
  done < <(cd "$repo" && find . -type f ! -name '*.heal' ! -name '*.1' -printf '%P\n' | sort)
  log "local $repo: $ok files verified, $made par2 sets created, status $STATUS"
}

# ------------------------------------------------------------------- b2 ----
b2_setup() {
  [[ -n "$RCLONE" ]] || die "rclone not installed"
  if [[ -n "${HOSTREPO_HEAL_ROOT:-}" ]]; then
    # Any rclone path laid out like the bucket (<root>/<HOST>, <root>/par2/<HOST>).
    REPO="$HOSTREPO_HEAL_ROOT/$1"; SETS="$HOSTREPO_HEAL_ROOT/par2/$1"; BUCKET=""
  else
    local bucket="${HOSTREPO_B2_BUCKET:-}"
    if [[ -z "$bucket" && "${RESTIC_REPOSITORY:-}" == b2:* ]]; then bucket="${RESTIC_REPOSITORY#b2:}"; bucket="${bucket%%:*}"; fi
    local acct="${B2_ACCOUNT_ID:-${MIRROR_B2_ACCOUNT:-}}" key="${B2_ACCOUNT_KEY:-${MIRROR_B2_KEY:-}}"
    [[ -n "$bucket" && -n "$acct" && -n "$key" ]] || die "need a bucket (HOSTREPO_B2_BUCKET or RESTIC_REPOSITORY) and a B2 key"
    export RCLONE_CONFIG=/dev/null RCLONE_CONFIG_HR_TYPE=b2 RCLONE_CONFIG_HR_ACCOUNT="$acct" RCLONE_CONFIG_HR_KEY="$key"
    REPO="hr:$bucket/$1"; SETS="hr:$bucket/par2/$1"; BUCKET="$bucket"
  fi
  REPO_REL="$1"; SETS_REL="par2/$1"
  WORK=$(mktemp -d); trap 'rm -rf "$WORK"' EXIT
}
b2_files() { "$RCLONE" lsf -R --files-only "$1" 2>/dev/null | sort; }
# put_to <bucket-relative base> <subdir or ""> <file>...: upload with the
# native API (hostrepo-b2put): B2 verifies each file's SHA-1 on receipt and
# stores its SHA-256 as file info. In test mode (HOSTREPO_HEAL_ROOT), a copy.
put_to() {
  local dir="$1${2:+/$2}"; shift 2
  if [[ -z "$BUCKET" ]]; then
    local f; for f in "$@"; do "$RCLONE" copyto -q "$f" "$HOSTREPO_HEAL_ROOT/$dir/$(basename -- "$f")" || return 2; done
  else
    "$(dirname "$0")/hostrepo-b2put" "$BUCKET" "$dir" "$@"
  fi
}

# SHA1SUMS in par2/<HOST>/ holds "<sha1>  <path>" for every repository file
# whose bytes were proven good (SHA-256 equal to its name). b2-lite compares
# it with the SHA-1 that B2 reports in a listing, with no download.
sums_get() { "$RCLONE" cat "$SETS/SHA1SUMS" 2>/dev/null | sort -k2 > "$WORK/sums" || : > "$WORK/sums"; }
sums_put() { sort -u -k2 "$WORK/sums" -o "$WORK/sums"; mkdir -p "$WORK/s"; cp "$WORK/sums" "$WORK/s/SHA1SUMS"; put_to "$SETS_REL" "" "$WORK/s/SHA1SUMS"; }

cmd_b2_sync() {
  b2_setup "${1:?HOST}"
  b2_files "$REPO" | while IFS= read -r r; do addressed "$r" && echo "$r"; done > "$WORK/files"
  b2_files "$SETS" | sed -n 's/\.par2$//p' | grep -v '\.vol[0-9+]*$' | sort > "$WORK/sets" || true
  sums_get; cp "$WORK/sums" "$WORK/sums.before"
  local new stale unsummed made=0 summed=0
  new=$(comm -23 "$WORK/files" "$WORK/sets"); stale=$(comm -13 "$WORK/files" "$WORK/sets")
  # Files that have a set but no SHA-1 yet: verify them once, record the SHA-1.
  unsummed=$(comm -12 "$WORK/files" "$WORK/sets" | comm -23 - <(awk '{print $2}' "$WORK/sums" | sort))
  while IFS= read -r rel; do
    [[ -n "$rel" ]] || continue
    [[ $DRY -eq 1 ]] && { log "dry-run: would create the set for $rel"; continue; }
    mkdir -p "$WORK/r/$(dirname "$rel")"
    "$RCLONE" copyto -q "$REPO/$rel" "$WORK/r/$rel" || { fail "$rel: download failed"; continue; }
    hash_ok "$WORK/r/$rel" "$rel" || { fail "$rel: DAMAGED in B2 (hash mismatch on first read); run b2-repair"; rm -f "$WORK/r/$rel"; continue; }
    par2_create "$WORK/r" "$rel" "$WORK/p" || { fail "$rel: par2 create failed"; continue; }
    put_to "$SETS_REL" "$(dirname "$rel")" "$WORK/p/$rel".par2 "$WORK/p/$rel".vol*.par2 || { fail "$rel: set upload failed"; continue; }
    printf '%s  %s\n' "$(sha1sum < "$WORK/r/$rel" | cut -c1-40)" "$rel" >> "$WORK/sums"
    rm -rf "${WORK:?}/r" "${WORK:?}/p"; made=$((made+1))
  done <<< "$new"
  while IFS= read -r rel; do
    [[ -n "$rel" && $DRY -eq 0 ]] || continue
    mkdir -p "$WORK/r/$(dirname "$rel")"
    "$RCLONE" copyto -q "$REPO/$rel" "$WORK/r/$rel" || { fail "$rel: download failed"; continue; }
    if hash_ok "$WORK/r/$rel" "$rel"; then
      printf '%s  %s\n' "$(sha1sum < "$WORK/r/$rel" | cut -c1-40)" "$rel" >> "$WORK/sums"; summed=$((summed+1))
    else fail "$rel: DAMAGED in B2; run b2-repair"; fi
    rm -rf "${WORK:?}/r"
  done <<< "$unsummed"
  local gone=0
  while IFS= read -r rel; do
    [[ -n "$rel" ]] || continue
    [[ $DRY -eq 1 ]] && { log "dry-run: would delete the set of pruned $rel"; continue; }
    "$RCLONE" delete -q "$SETS/$(dirname "$rel")" --include "${rel##*/}*.par2" && gone=$((gone+1)) || fail "$rel: could not delete stale set"
  done <<< "$stale"
  # Keep SHA-1s only for files that still exist.
  awk 'NR==FNR {keep[$1]=1; next} ($2 in keep)' "$WORK/files" "$WORK/sums" > "$WORK/sums.new"; mv "$WORK/sums.new" "$WORK/sums"
  if [[ $DRY -eq 0 ]] && ! cmp -s <(sort -u -k2 "$WORK/sums") <(sort -u -k2 "$WORK/sums.before"); then sums_put || fail "could not store SHA1SUMS"; fi
  log "b2-sync $1: $(wc -l < "$WORK/files") files, $made sets created, $summed SHA-1s recorded, $gone stale sets deleted, status $STATUS"
}

# b2-lite <HOST> [--repair]: compare the SHA-1 B2 reports for every repository
# file with SHA1SUMS. No file is downloaded unless it fails, and then only to
# repair it (--repair).
cmd_b2_lite() {
  local host="${1:?HOST}" repair=0; [[ "${2:-}" == --repair ]] && repair=1
  b2_setup "$host"; sums_get
  "$RCLONE" lsjson -R --files-only --hash --hash-type sha1 "$REPO" 2>/dev/null |
    jq -r '.[] | [.Path, (.Hashes.sha1 // "")] | @tsv' | sort > "$WORK/listed" || die "cannot list $REPO"
  local n=0 ok=0 bad=0 unver=0 nohash=0 missing=0 rel sha want ids=()
  declare -A stored=(); while read -r want rel; do stored[$rel]=$want; done < "$WORK/sums"
  declare -A seen=()
  while IFS=$'\t' read -r rel sha; do
    addressed "$rel" || continue
    n=$((n+1)); seen[$rel]=1
    want="${stored[$rel]:-}"
    if [[ -z "$want" ]]; then unver=$((unver+1))
    elif [[ -z "$sha" ]]; then nohash=$((nohash+1))
    elif [[ "$sha" == "$want" ]]; then ok=$((ok+1))
    else bad=$((bad+1)); log "DAMAGED in B2: $rel (SHA-1 $sha, recorded $want)"; ids+=("${rel##*/}"); fi
  done < "$WORK/listed"
  for rel in "${!stored[@]}"; do
    [[ -n "${seen[$rel]:-}" ]] || { missing=$((missing+1)); fail "$rel: MISSING from B2 (recorded as verified); restore it from the EU replica or a peer mirror"; }
  done
  log "b2-lite $host: $n files, $ok match, $bad mismatched, $unver not yet verified, $nohash without a SHA-1 in B2, $missing missing"
  if [[ $bad -gt 0 ]]; then
    if [[ $repair -eq 1 ]]; then cmd_b2_repair "$host" "${ids[@]}"; else STATUS=2; fi
  fi
}

cmd_b2_repair() {
  local host="${1:?HOST}"; shift; b2_setup "$host"
  b2_files "$REPO" > "$WORK/files"
  local id rel
  for id in "$@"; do
    rel=$(grep -E "(^|/)$id\$" "$WORK/files" | head -1 || true)
    [[ -n "$rel" ]] || continue
    mkdir -p "$WORK/r/$(dirname "$rel")" "$WORK/p/$(dirname "$rel")"
    "$RCLONE" copyto -q "$REPO/$rel" "$WORK/r/$rel" || { fail "$rel: download failed"; continue; }
    if hash_ok "$WORK/r/$rel" "$rel"; then log "$rel: intact in B2"; continue; fi
    log "DAMAGED in B2: $rel"
    [[ $DRY -eq 1 ]] && { STATUS=2; continue; }
    "$RCLONE" copy -q "$SETS/$(dirname "$rel")" "$WORK/p/$(dirname "$rel")" --include "${rel##*/}*.par2" || true
    if [[ -f "$WORK/p/$rel.par2" ]] && par2_repair "$WORK/r" "$rel" "$WORK/p" && hash_ok "$WORK/r/$rel" "$rel"; then
      put_to "$REPO_REL" "$(dirname "$rel")" "$WORK/r/$rel" && warn "$rel: repaired from par2 and re-uploaded" || fail "$rel: repaired but upload failed"
    else
      fail "$rel: no usable par2 set; restore it from the EU replica or a peer mirror"
    fi
  done
}

# lite-drill <HOST> [--mirror DIR] [--json FILE]
# The nightly drill without downloads: every copy of HOST's repository against
# SHA1SUMS. The US bucket by listing; a local mirror by hashing its files; the
# EU replica by listing when HOSTREPO_EU_ACCOUNT/KEY/BUCKET are set. Reports,
# never repairs (the backup and the mirror repair their own copies).
cmd_lite_drill() {
  local host="${1:?HOST}" mirror="" json=""; shift
  while [[ $# -gt 0 ]]; do case "$1" in --mirror) mirror="$2"; shift ;; --json) json="$2"; shift ;; *) die "unknown option $1" ;; esac; shift; done
  b2_setup "$host"; sums_get
  local total; total=$(wc -l < "$WORK/sums")
  [[ $total -gt 0 ]] || { log "lite-drill $host: no SHA1SUMS yet; nothing to compare"; STATUS=1; }
  # compare <name> <file of "path<TAB>sha1">: counts against SHA1SUMS
  compare() {
    awk -F'\t' -v name="$1" 'NR==FNR {split($0, a, "  "); want[a[2]]=a[1]; next}
      { have[$1]=$2 }
      END { m=0; b=0; x=0; for (p in want) { if (!(p in have)) x++; else if (have[p]==want[p]) m++; else { b++; print "MISMATCH " name ": " p > "/dev/stderr" } }
            printf "%d %d %d\n", m, b, x }' "$WORK/sums" "$2"
  }
  local us eu="" mi=""
  "$RCLONE" lsjson -R --files-only --hash --hash-type sha1 "$REPO" 2>/dev/null | jq -r '.[] | [.Path, (.Hashes.sha1 // "")] | @tsv' > "$WORK/us.tsv" || die "cannot list $REPO"
  us=$(compare us "$WORK/us.tsv")
  if [[ -n "$mirror" && -d "$mirror" ]]; then
    (cd "$mirror" && awk '{print $2}' "$WORK/sums" | while IFS= read -r p; do if [[ -f "$p" ]]; then printf '%s\t%s\n' "$p" "$(sha1sum < "$p" | cut -c1-40)"; fi; done) > "$WORK/mirror.tsv"
    mi=$(compare mirror "$WORK/mirror.tsv")
  fi
  if [[ -n "${HOSTREPO_EU_ACCOUNT:-}" && -n "${HOSTREPO_EU_KEY:-}" && -n "${HOSTREPO_EU_BUCKET:-}" ]]; then
    RCLONE_CONFIG_EU_TYPE=b2 RCLONE_CONFIG_EU_ACCOUNT="$HOSTREPO_EU_ACCOUNT" RCLONE_CONFIG_EU_KEY="$HOSTREPO_EU_KEY" \
      "$RCLONE" lsjson -R --files-only --hash --hash-type sha1 "eu:$HOSTREPO_EU_BUCKET/$host" 2>/dev/null |
      jq -r '.[] | [.Path, (.Hashes.sha1 // "")] | @tsv' > "$WORK/eu.tsv" && eu=$(compare eu "$WORK/eu.tsv")
  fi
  local bad=0 n
  for n in "$us" "$mi" "$eu"; do [[ -z "$n" ]] && continue; set -- $n; bad=$((bad + $2 + $3)); done
  [[ $bad -gt 0 ]] && STATUS=2
  log "lite-drill $host: $total recorded files; us $us${mi:+; mirror $mi}${eu:+; eu $eu} (match mismatched missing); status $STATUS"
  if [[ -n "$json" ]]; then
    jq -n --arg us "$us" --arg mi "$mi" --arg eu "$eu" --argjson total "$total" --argjson rc "$STATUS" '
      def c($s): if $s == "" then null else ($s | split(" ") | map(tonumber) | {match: .[0], mismatched: .[1], missing: .[2]}) end;
      {kind: "litedrill", result: $rc, checked: (now | todate), files: $total, us: c($us), mirror: c($mi), eu: c($eu)}' > "$json"
  fi
}

cmd="${1:-}"; shift || true
args=(); for a in "$@"; do [[ "$a" == --dry-run ]] && DRY=1 || args+=("$a"); done
case "$cmd" in
  local)     cmd_local "${args[@]}" ;;
  b2-sync)   cmd_b2_sync "${args[@]}" ;;
  b2-repair) cmd_b2_repair "${args[@]}" ;;
  b2-lite)   cmd_b2_lite "${args[@]}" ;;
  lite-drill) cmd_lite_drill "${args[@]}" ;;
  *) sed -n '2,40p' "$0" >&2; exit 2 ;;
esac
exit $STATUS
