#!/usr/bin/env bash
# Forensic recovery of deleted ~/.claude/ data from a btrfs partition.
#
# Two read-only strategies, applied in order against the SAME source (image
# preferred, live device discouraged):
#
#  (A) btrfs tree-walk: btrfs is copy-on-write, so older root trees may
#      persist on disk until allocator GC overwrites them. We enumerate them
#      with `btrfs-find-root`, then use `btrfs restore` (which opens the
#      source read-only) with the prefix-form `--path-regex` documented in
#      `man btrfs-restore` to extract just the ~/.claude subtree from each
#      candidate.
#
#  (B) raw-device carving: stream the source in chunks through tr/grep to
#      find ERE-matching JSON fragments characteristic of Claude Code logs.
#      Memory-bounded (lines kept short via tr); output capped at 1G.
#
# DEFAULT MODE requires --image PATH (a `dd` image of the partition). Live
# devices are accepted only with --device PATH and either --force-live or
# unmounted, because `btrfs restore` against a mounted FS can return
# silently corrupt output spliced from concurrent kernel writes.
#
# Usage:
#   sudo bash scripts/recover-btrfs --image /mnt/ext/home.img
#   sudo bash scripts/recover-btrfs --device /dev/nvme0n1p6 --force-live
#   sudo bash scripts/recover-btrfs --image ... --dry-run
#
# Other flags:
#   --out PATH         Output base (default /var/tmp/btrfs-recover).
#                      Must be under /var/tmp /mnt /media /run/media.
#   --max-trees N      Cap candidate-tree iteration (default 50).
#   --path-regex RX    Override path-regex (default: nested-prefix form for
#                      ~/.claude only -- does NOT match .claude-personal etc.).
#   --max-output BYTES Output cap for carving (default 1073741824).
#   --skip-trees       Skip phase A.
#   --skip-carve       Skip phase B.
#   --dry-run          Print what would run; do not execute.

set -euo pipefail
IFS=$'\n\t'

# ─────────────────────────────────────────────────────────────────────────────
# Defaults
# ─────────────────────────────────────────────────────────────────────────────

SOURCE=""
SOURCE_KIND=""           # "image" or "device"
OUT_BASE=""
MAX_TREES=50
DRY_RUN=0
FORCE_LIVE=0
SKIP_TREES=0
SKIP_CARVE=0
MAX_OUTPUT_BYTES=$((1024 * 1024 * 1024))    # 1G default cap

# Nested-prefix path regex required by `btrfs restore --path-regex`. The
# walker tests each level from `/` downward; a flat regex like
# `^/home/m/\.claude($|/)` would fail at `/home` and stop traversal.
# The form below allows traversing `/`, `/home`, `/home/m`, `/home/m/.claude`,
# and any descendant, but NOT siblings like `.claude-personal`.
PATH_REGEX_DEFAULT='^(|/(|home(|/m(|/\.claude(|/.*)))))$'
PATH_REGEX="$PATH_REGEX_DEFAULT"

# Carve patterns. ERE-correct (bare {N}, not BRE \{N\}). Matching tries to
# find any line containing the substring; raw-device "lines" are accidental
# (we synthesize them via tr below) but JSONL records on disk preserve
# their internal printable bytes intact, so each pattern targets a token
# that real Claude Code records contain.
#
# Source-of-truth fields (from src/claude_timeline/build/{messages,bash,
# assistant,tools}.py and the per-session JSONL format):
#   history.jsonl: display, timestamp, project, sessionId, pastedContents
#   session UUID jsonl: type, sessionId, parentUuid, message.content[].text,
#     message.content[].input.command (Bash), cwd, gitBranch, agentId
CARVE_PATTERNS=(
  '"display":"'
  '"type":"(user|assistant|system)"'
  '"sessionId":"[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}"'
  '"name":"Bash"'
  '"input":\{"command":"'
  '"parentUuid":"[a-f0-9]{8}-[a-f0-9]{4}-'
  '"cwd":"/home/m'
  '"gitBranch":"'
  '"agentId":"'
)

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

die()  { printf 'ERROR: %s\n' "$*" >&2; exit 1; }
log()  { printf '%s\n' "$*" | tee -a "${LOG:-/dev/null}"; }

usage() {
  sed -n '/^# Usage:/,/^$/p' "$0" | sed 's/^# \{0,1\}//' >&2
}

# ─────────────────────────────────────────────────────────────────────────────
# Argument parsing
# ─────────────────────────────────────────────────────────────────────────────

while [[ $# -gt 0 ]]; do
  case "$1" in
    --image)        SOURCE_KIND=image;  SOURCE="${2:-}"; shift 2 ;;
    --device)       SOURCE_KIND=device; SOURCE="${2:-}"; shift 2 ;;
    --out)          OUT_BASE="${2:-}";  shift 2 ;;
    --max-trees)    MAX_TREES="${2:-}"; shift 2 ;;
    --path-regex)   PATH_REGEX="${2:-}"; shift 2 ;;
    --max-output)   MAX_OUTPUT_BYTES="${2:-}"; shift 2 ;;
    --skip-trees)   SKIP_TREES=1; shift ;;
    --skip-carve)   SKIP_CARVE=1; shift ;;
    --dry-run)      DRY_RUN=1; shift ;;
    --force-live)   FORCE_LIVE=1; shift ;;
    -h|--help)      usage; exit 0 ;;
    *) printf 'unknown arg: %s\n' "$1" >&2; usage; exit 1 ;;
  esac
done

# ─────────────────────────────────────────────────────────────────────────────
# Pre-flight validation
# ─────────────────────────────────────────────────────────────────────────────

[[ $EUID -eq 0 ]] || die "must run as root (try: sudo bash $0 ...)"
[[ -n "$SOURCE" ]] || die "specify --image PATH or --device PATH"
[[ -n "$SOURCE_KIND" ]] || die "specify --image or --device"
[[ "$MAX_TREES" =~ ^[0-9]+$ ]] || die "--max-trees must be integer"
[[ "$MAX_OUTPUT_BYTES" =~ ^[0-9]+$ ]] || die "--max-output must be integer bytes"

case "$SOURCE_KIND" in
  image)
    [[ -f "$SOURCE" ]] || die "image file not found: $SOURCE"
    ;;
  device)
    [[ -b "$SOURCE" ]] || die "not a block device: $SOURCE"
    ;;
esac

# Verify FS type
fstype=$(blkid -s TYPE -o value -- "$SOURCE" 2>/dev/null || true)
[[ "$fstype" == btrfs ]] || die "source fstype='$fstype', expected btrfs (use blkid to verify)"

# Mounted-rw guard for live devices
if [[ "$SOURCE_KIND" == device ]]; then
  mounted_opts=$(findmnt -no OPTIONS --source "$SOURCE" 2>/dev/null | head -1 || true)
  if [[ -n "$mounted_opts" ]]; then
    if grep -qw rw <<<"$mounted_opts"; then
      if [[ $FORCE_LIVE -ne 1 ]]; then
        die "device $SOURCE is mounted rw; results would be unreliable. Either:
  - dd it to an image and re-run with --image PATH, or
  - boot from a live USB and run there, or
  - pass --force-live (not recommended; concurrent kernel writes can yield
    silently corrupt 'recovered' files spliced from in-flight metadata)"
      fi
      printf '\nWARNING: --force-live mode against mounted rw device.\n' >&2
      printf 'Recovered files may be silently corrupt. Sleeping 10s -- Ctrl-C to abort.\n\n' >&2
      sleep 10
    fi
  fi
fi

# Output base validation
OUT_BASE="${OUT_BASE:-/var/tmp/btrfs-recover}"
OUT_BASE="$(realpath -m -- "$OUT_BASE")"
case "$OUT_BASE" in
  /var/tmp/*|/mnt/*|/media/*|/run/media/*) ;;
  *) die "OUT_BASE must be under /var/tmp /mnt /media /run/media; got: $OUT_BASE" ;;
esac

# Free-space sanity (need at least 5G)
parent="$OUT_BASE"
while [[ ! -d "$parent" ]]; do parent="$(dirname "$parent")"; done
avail_g=$(df -BG --output=avail "$parent" 2>/dev/null | tail -1 | tr -dc '0-9')
[[ "${avail_g:-0}" -ge 5 ]] || die "Need ≥ 5G free in $parent; have ${avail_g:-0}G"

# Output directories (timestamped to avoid mixing with prior runs)
TS=$(date -u +%Y%m%dT%H%M%SZ)
LOG_DIR="$OUT_BASE/logs"
TREES_DIR="$OUT_BASE/trees-$TS"
CARVED_DIR="$OUT_BASE/carved-$TS"
install -m 700 -d "$OUT_BASE" "$LOG_DIR" "$TREES_DIR" "$CARVED_DIR"
LOG="$LOG_DIR/recover-$TS.log"
ln -sfn "recover-$TS.log" "$LOG_DIR/latest.log"

# Cleanup hook: chown to the invoking user, restorecon, prune empty dirs
cleanup() {
  local rc=$?
  log "[trap] exiting rc=$rc"
  if [[ -n "${SUDO_USER:-}" ]]; then
    chown -R "$SUDO_USER":"$SUDO_USER" "$OUT_BASE" 2>>"$LOG" || true
  else
    log "[trap] WARNING: SUDO_USER unset (script run via su or root shell?)."
    log "        Output files in $OUT_BASE remain owned by root."
    log "        Run as a non-root user: chown -R \$USER $OUT_BASE"
  fi
  command -v restorecon >/dev/null 2>&1 && restorecon -R "$OUT_BASE" 2>>"$LOG" || true
  # Prune empty timestamped dirs from this run if we produced nothing.
  [[ -d "$TREES_DIR"  && -z "$(ls -A "$TREES_DIR"  2>/dev/null)" ]] && rmdir "$TREES_DIR"  2>/dev/null || true
  [[ -d "$CARVED_DIR" && -z "$(ls -A "$CARVED_DIR" 2>/dev/null)" ]] && rmdir "$CARVED_DIR" 2>/dev/null || true
  exit $rc
}
trap cleanup EXIT INT TERM

# Lower priority: this is a read-heavy 30-90 min scan against the active
# disk; without ionice the desktop becomes sluggish. Best-effort -- failures
# are non-fatal (e.g., ionice missing, kernel without scheduler classes).
ionice -c 3 -p $$ 2>/dev/null || true
renice -n 19 $$ >/dev/null 2>&1 || true

# ─────────────────────────────────────────────────────────────────────────────
# Banner
# ─────────────────────────────────────────────────────────────────────────────

log "[recover] timestamp=$TS"
log "[recover] source=$SOURCE kind=$SOURCE_KIND"
log "[recover] out=$OUT_BASE max-trees=$MAX_TREES dry-run=$DRY_RUN force-live=$FORCE_LIVE"
log "[recover] btrfs=$(btrfs --version 2>/dev/null | head -1) kernel=$(uname -r)"
log "[recover] path-regex=$PATH_REGEX"
log "[recover] carve patterns: ${#CARVE_PATTERNS[@]}"
for p in "${CARVE_PATTERNS[@]}"; do log "  - $p"; done

# Warn if output and source are on the same device. For Phase A this matters
# (every output write reduces what older trees can recover); for Phase B
# (carving past the deletion event) it doesn't, but flag for awareness.
if [[ "$SOURCE_KIND" == device ]]; then
  out_dev=$(df --output=source "$OUT_BASE" 2>/dev/null | tail -1 || true)
  if [[ "$out_dev" == "$SOURCE" ]]; then
    log "[recover] NOTE: output is on the same device as source ($SOURCE)."
    log "          Phase B is unaffected. For Phase A use --out on another disk."
  fi
fi

if [[ $DRY_RUN -eq 1 ]]; then
  log "===== DRY RUN ====="
  log "[A] would call: btrfs-find-root $SOURCE"
  log "[A] then for each tree, capped at $MAX_TREES:"
  log "    btrfs restore -t <bytenr> -i -v --path-regex '$PATH_REGEX' $SOURCE <out>"
  log "[B] would scan $SOURCE for ${#CARVE_PATTERNS[@]} ERE patterns:"
  for p in "${CARVE_PATTERNS[@]}"; do log "      $p"; done
  log "[B] output capped at $MAX_OUTPUT_BYTES bytes"
  log "===== end dry run ====="
  exit 0
fi

# ─────────────────────────────────────────────────────────────────────────────
# (A) btrfs tree-walk
# ─────────────────────────────────────────────────────────────────────────────

if [[ $SKIP_TREES -eq 0 ]]; then
  log "===== (A) btrfs tree-walk ====="
  if ! command -v btrfs-find-root >/dev/null 2>&1; then
    log "[A] btrfs-find-root missing (install btrfs-progs); skipping (A)."
  else
    ROOTS_FILE="$LOG_DIR/find-root-$TS.txt"
    log "[A] enumerating older trees..."
    btrfs-find-root "$SOURCE" >"$ROOTS_FILE" 2>>"$LOG" || true
    log "[A] $(wc -l < "$ROOTS_FILE") line(s) from btrfs-find-root"

    # Strict parser: only pick the bytenr field. Two known output forms across
    # btrfs-progs versions:
    #   "Found tree root at N gen M level L"   -> $5
    #   "Well block N(gen: M level: L)"        -> strip "(...)" from $3
    # Lower bound 1048576 (1MB) defensively rejects any leaked tree-level
    # numbers (which are small ints) -- real btrfs metadata bytenrs are
    # always well above 1MB. Without this, a parser glitch could feed
    # bytenr=1 to btrfs restore, wasting a 10-min timeout per bad value.
    mapfile -t CANDIDATES < <(
      awk '
        /^Found tree root at[[:space:]]/ {
          if ($5 ~ /^[0-9]+$/ && $5 + 0 >= 1048576) print $5
        }
        /^Well block[[:space:]]/ {
          b = $3; sub(/\(.*$/, "", b)
          if (b ~ /^[0-9]+$/ && b + 0 >= 1048576) print b
        }
      ' "$ROOTS_FILE" | awk '!seen[$0]++' | head -n "$MAX_TREES"
    )

    if [[ ${#CANDIDATES[@]} -eq 0 ]]; then
      log "[A] no candidate trees parsed from btrfs-find-root output."
      log "    The on-disk tree history may be fully overwritten, or"
      log "    btrfs-find-root output format is unrecognized. Inspect $ROOTS_FILE."
    else
      log "[A] trying ${#CANDIDATES[@]} candidate(s) (capped at $MAX_TREES)..."
      hits=0
      for i in "${!CANDIDATES[@]}"; do
        bytenr="${CANDIDATES[$i]}"
        out=$(mktemp -d -p "$TREES_DIR" "gen-${bytenr}.XXXXXX")
        log "[A] [$((i+1))/${#CANDIDATES[@]}] tree@$bytenr -> $out"
        if timeout 600 btrfs restore -t "$bytenr" -i -v \
             --path-regex "$PATH_REGEX" "$SOURCE" "$out" >>"$LOG" 2>&1; then
          # Tighten hit count to JSONL-shaped files; restored directory
          # stubs / lockfiles / .tmp artifacts are not useful and would
          # inflate the count.
          n=$(find "$out" -type f -size +0c -name '*.jsonl' 2>/dev/null | wc -l)
          if [[ $n -gt 0 ]]; then
            sz=$(du -sh "$out" 2>/dev/null | cut -f1)
            log "    +$n file(s), $sz"
            hits=$((hits + 1))
          else
            rm -rf "$out"
          fi
        else
          log "    (failed or 10-min timeout)"
          rm -rf "$out"
        fi
      done
      log "[A] $hits tree(s) yielded files."
    fi
  fi
  log
fi

# ─────────────────────────────────────────────────────────────────────────────
# (B) raw-device chunked carving
# ─────────────────────────────────────────────────────────────────────────────

if [[ $SKIP_CARVE -eq 0 ]]; then
  log "===== (B) chunked raw carving ====="
  CARVED_FILE="$CARVED_DIR/claude-candidates.jsonl"
  : > "$CARVED_FILE"

  total_size=$(blockdev --getsize64 "$SOURCE" 2>/dev/null || stat -c %s -- "$SOURCE")
  total_g=$((total_size / 1024 / 1024 / 1024))
  log "[B] source size: ${total_g}G; output cap: $MAX_OUTPUT_BYTES bytes"

  # Combine all patterns into one ERE alternation.
  COMBINED="$(IFS='|'; printf '%s' "${CARVE_PATTERNS[*]}")"
  # Line-mode grep -aE (no -o, no context-extension regex). The earlier
  # `-o` + `.{0,4096}` form caused catastrophic backtracking in GNU grep
  # 3.12 (~25 KB/s on history.jsonl, scaling to ~50 hours for 243 GB).
  # Line-mode bails on first match per line: ~300 MB/s, ~14 min for 243 GB.
  # The `tr` step below replaces control bytes with newlines, so each
  # "line" emerging from tr is bounded by the longest run of non-control
  # bytes on disk -- in practice MB-scale, well within grep's per-line
  # buffer. JSONL records (purely printable except '\n') flow through
  # intact as single tr-lines and are emitted whole when they match.
  GREP_PATTERN="(${COMBINED})"

  # Memory bound: tr replaces only true control bytes (NUL through BS, VT
  # through US, DEL) with newlines. This bounds grep's per-line buffer to
  # the longest run of non-control bytes on disk (a few MB max). Critically,
  # this PRESERVES UTF-8 multibyte sequences (0x80-0xFF), so emoji, smart
  # quotes, accented characters in real Claude data survive intact. The
  # earlier `tr -c '[:print:]\t\n'` form would shred all non-ASCII text
  # because UTF-8 continuation bytes are non-printable in the C locale.
  #
  # iflag=direct bypasses page cache (avoids polluting kernel cache with
  # 243G of partition content). On mounted btrfs it can fail EINVAL; we
  # fall back to buffered reads if direct I/O isn't accepted.
  #
  # head -c caps total output bytes; pipeline closes cleanly via SIGPIPE
  # after that.
  log "[B] streaming source through tr | grep | head ..."
  IFLAG="iflag=direct"
  # Probe direct I/O first with a tiny read -- if it fails we drop iflag.
  if ! dd if="$SOURCE" bs=8M iflag=direct count=1 of=/dev/null status=none 2>/dev/null; then
    log "[B] iflag=direct rejected on $SOURCE; falling back to buffered reads"
    IFLAG=""
  fi
  set +e
  # dd progress: tee to terminal AND log so the user sees live status during
  # the carve run.
  dd if="$SOURCE" bs=8M $IFLAG status=progress 2> >(tee -a "$LOG" >&2) \
    | LC_ALL=C tr '\000-\010\013-\037\177' '\n' \
    | LC_ALL=C grep -aE --binary-files=text --line-buffered "$GREP_PATTERN" 2>>"$LOG" \
    | head -c "$MAX_OUTPUT_BYTES" \
    > "$CARVED_FILE" 2>>"$LOG"
  carve_rc=$?
  set -e
  # rc semantics under set -o pipefail (still active during set +e):
  #   141 = SIGPIPE from head closing pipe at output cap (success, capped)
  #   1   = grep found zero matches across full source (success, no hits)
  #   0   = pipeline finished cleanly with at least one match
  #   other = real failure (dd I/O error, tr fail, etc.)
  if [[ $carve_rc -eq 0 || $carve_rc -eq 141 ]]; then
    :  # success
  elif [[ $carve_rc -eq 1 && ! -s "$CARVED_FILE" ]]; then
    log "[B] grep found no matches across full source (rc=1); carved file is empty"
  else
    log "[B] WARNING: carve pipeline exited rc=$carve_rc -- output may be incomplete"
  fi

  if [[ -s "$CARVED_FILE" ]]; then
    # Drop torn final line: `head -c` truncates at byte boundary, so the
    # last line may be a fragment with no trailing newline. Otherwise that
    # fragment confuses downstream JSON parsers.
    if [[ "$(tail -c1 "$CARVED_FILE" | od -An -tx1 | tr -d ' ')" != "0a" ]]; then
      log "[B] dropping torn final line (truncated mid-record at output cap)"
      sed -i '$d' "$CARVED_FILE"
    fi
    log "[B] de-duplicating..."
    sort -u "$CARVED_FILE" -o "$CARVED_FILE" 2>>"$LOG" || true
    n=$(wc -l < "$CARVED_FILE")
    sz=$(du -sh "$CARVED_FILE" 2>/dev/null | cut -f1)
    log "[B] $n unique candidate line(s), $sz -> $CARVED_FILE"
  else
    log "[B] no matches; carved file is empty."
  fi
fi
log

# ─────────────────────────────────────────────────────────────────────────────
# Summary
# ─────────────────────────────────────────────────────────────────────────────

log "===== summary ====="
log "trees output:  $TREES_DIR"
log "carved output: $CARVED_DIR"
log "log:           $LOG"

if [[ -d "$TREES_DIR" ]]; then
  echo
  echo "Surviving files from (A) (first 30):"
  find "$TREES_DIR" -type f 2>/dev/null | head -30 || true
fi

CARVED_FILE="$CARVED_DIR/claude-candidates.jsonl"
if [[ -s "$CARVED_FILE" ]]; then
  echo
  echo "Sample carved lines from (B) (first 5, control chars stripped):"
  head -5 "$CARVED_FILE" | LC_ALL=C tr -d '\000-\010\013-\037\177' | cut -b1-200
fi
