#!/usr/bin/env bash
# mempalace-backup — consistent, verified backups of a file-backed MemPalace palace.
#
# WHY THIS EXISTS
#   The fleet primary (RFC-001) had no backups at all. A palace is not one file:
#   it is two SQLite databases (one of them in WAL mode), a set of chromadb HNSW
#   segment directories, and a few JSON sidecars. `cp -r` on a live palace can
#   capture a torn WAL and an index that disagrees with the sqlite.
#
# TWO MODES
#   hot  (default)  Zero downtime. Both SQLite files are copied with the sqlite
#                   ONLINE BACKUP API (a transactionally consistent snapshot,
#                   WAL included). The HNSW segment dirs are rsynced and may be
#                   slightly skewed relative to the sqlite — which is RECOVERABLE,
#                   because the sqlite is the source of truth and the index can be
#                   rebuilt: `mempalace repair --mode from-sqlite`.
#   cold (--cold)   Stops the server, copies everything byte-for-byte, restarts.
#                   Guaranteed self-consistent including the index. The restart is
#                   trapped so an error still brings the server back.
#
# BONUS: restoring the hot backup via `repair --mode from-sqlite --archive-existing`
#   re-CREATES the collections, which is also the only path that applies the modern
#   chromadb HNSW defaults (batch_size=100 / sync_threshold=1000). Palaces created
#   under mempalace <=3.6 are stuck on 2/2 and sync the index every 2 records.
#
# NOTE: synlig has no sqlite3 CLI, so every SQLite operation here goes through
#   python3's stdlib sqlite3 module. Do not "simplify" this to `VACUUM INTO`.
set -euo pipefail

PALACE="${MEMPALACE_PALACE_PATH:-$HOME/.mempalace/palace}"
DEST="${MEMPALACE_BACKUP_DIR:-$HOME/backups/mempalace}"
UNIT="mempalace-serve"
KEEP=7
COLD=0
DRY=0
VERIFY=""

usage() {
  cat <<'USAGE'
usage: mempalace-backup [--cold] [--dest DIR] [--keep N] [--palace DIR] [--dry-run]
       mempalace-backup --verify BACKUP_DIR

  --cold        quiesce the server (stop/copy/start) for a byte-consistent copy
  --dest DIR    backup root (default: $HOME/backups/mempalace)
  --keep N      how many backups to retain (default: 7; 0 = keep all)
  --palace DIR  palace to back up (default: $MEMPALACE_PALACE_PATH or ~/.mempalace/palace)
  --dry-run     report what would happen, touch nothing
  --verify DIR  integrity-check an existing backup and print its counts
USAGE
}

while [ $# -gt 0 ]; do
  case "$1" in
    --cold) COLD=1 ;;
    --dest) DEST="$2"; shift ;;
    --keep) KEEP="$2"; shift ;;
    --palace) PALACE="$2"; shift ;;
    --dry-run) DRY=1 ;;
    --verify) VERIFY="$2"; shift ;;
    -h|--help) usage; exit 0 ;;
    *) echo "mempalace-backup: unknown option $1" >&2; usage >&2; exit 2 ;;
  esac
  shift
done

log() { printf '%s %s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)" "$*"; }
die() { log "ERROR: $*"; exit 1; }

# ── shared python helpers ────────────────────────────────────────────────────
# check_db: PRAGMA quick_check + row counts, on a copy. Prints JSON.
py_check() {
  python3 - "$1" <<'PY'
import json, os, sqlite3, sys
path = sys.argv[1]
out = {"file": os.path.basename(path), "bytes": os.path.getsize(path)}
try:
    c = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    out["quick_check"] = c.execute("PRAGMA quick_check").fetchone()[0]
    tables = [r[0] for r in c.execute(
        "SELECT name FROM sqlite_master WHERE type='table'")]
    out["tables"] = len(tables)
    counts = {}
    if "embedding_metadata" in tables:
        counts["drawers"] = c.execute(
            "SELECT count(*) FROM embedding_metadata "
            "WHERE key='chroma:document'").fetchone()[0]
    for t in ("facts", "triples", "entities"):
        if t in tables:
            counts[t] = c.execute(f"SELECT count(*) FROM {t}").fetchone()[0]
    out["counts"] = counts
    c.close()
except Exception as exc:  # noqa: BLE001 — report, never mask
    out["error"] = repr(exc)
print(json.dumps(out))
PY
}

# ── --verify: inspect an existing backup, then exit ──────────────────────────
if [ -n "$VERIFY" ]; then
  [ -d "$VERIFY" ] || die "no such backup dir: $VERIFY"
  log "verifying $VERIFY"
  rc=0
  for db in "$VERIFY"/chroma.sqlite3 "$VERIFY"/knowledge_graph.sqlite3; do
    [ -f "$db" ] || { log "  MISSING $(basename "$db")"; rc=1; continue; }
    res=$(py_check "$db")
    log "  $res"
    printf '%s' "$res" | grep -q '"quick_check": "ok"' || rc=1
  done
  if [ -f "$VERIFY/MANIFEST.json" ]; then
    python3 - "$VERIFY/MANIFEST.json" <<'PY'
import json, sys
m = json.load(open(sys.argv[1]))
print("  manifest: created={created_utc} mode={mode} bytes={bytes_total} "
      "mempalace={mempalace_version}".format(**m))
print("  index   :", m.get("index_consistency", "?"))
PY
  fi
  [ "$rc" -eq 0 ] && log "VERIFY OK" || log "VERIFY FAILED"
  exit "$rc"
fi

# ── preflight ───────────────────────────────────────────────────────────────
[ -d "$PALACE" ] || die "palace not found: $PALACE"
[ -f "$PALACE/chroma.sqlite3" ] || die "no chroma.sqlite3 in $PALACE — wrong dir?"
command -v rsync >/dev/null || die "rsync not found"

STAMP="$(date -u +%Y%m%dT%H%M%SZ)"
MODE=$([ "$COLD" -eq 1 ] && echo cold || echo hot)
OUT="$DEST/mp-$STAMP-$MODE"
TMP="$OUT.incomplete"

log "palace : $PALACE ($(du -sh "$PALACE" | cut -f1))"
log "dest   : $OUT"
log "mode   : $MODE   keep: $KEEP"

if [ "$DRY" -eq 1 ]; then
  log "dry-run: would create $OUT"
  log "dry-run: would copy chroma.sqlite3 + knowledge_graph.sqlite3 via sqlite backup API"
  log "dry-run: would rsync the palace tree ($(find "$PALACE" -mindepth 1 -maxdepth 1 | wc -l) top-level entries, hidden included) minus SQLite"
  [ "$COLD" -eq 1 ] && log "dry-run: would stop/start $UNIT"
  exit 0
fi

mkdir -p "$TMP"

# ── cold mode: quiesce, with a trap so the server always comes back ─────────
STOPPED=0
restart_if_stopped() {
  if [ "$STOPPED" -eq 1 ]; then
    log "restarting $UNIT"
    systemctl --user start "$UNIT" || log "WARNING: failed to start $UNIT — CHECK MANUALLY"
    STOPPED=0
  fi
}
trap restart_if_stopped EXIT INT TERM

if [ "$COLD" -eq 1 ]; then
  log "stopping $UNIT (quiesce)"
  systemctl --user stop "$UNIT"
  STOPPED=1
  # Give an in-flight mine a moment to release the writer lease.
  sleep 2
fi

# ── the copy ────────────────────────────────────────────────────────────────
# Order matters: rsync the tree FIRST, then write the two SQLite files as
# consistent snapshots, so the databases are the freshest thing in the backup.
#
# rsync takes everything EXCEPT the SQLite files (and their -wal/-shm): the HNSW
# segment dirs, the JSON sidecars, and hidden dirs such as .mempalace/origin.json
# — a `"$PALACE"/*/` glob silently skips those, which is exactly the bug this
# structure removes. The payoff: a backup dir IS a palace image, so restore is a
# copy rather than a procedure.
log "rsyncing palace tree (excluding SQLite)"
rsync -a \
  --exclude 'chroma.sqlite3'  --exclude 'chroma.sqlite3-*' \
  --exclude 'knowledge_graph.sqlite3' --exclude 'knowledge_graph.sqlite3-*' \
  "$PALACE"/ "$TMP"/

# SQLite: online backup API. Consistent even mid-write, and WAL-aware — which
# matters because knowledge_graph.sqlite3 runs in WAL mode (-wal/-shm present).
for db in chroma.sqlite3 knowledge_graph.sqlite3; do
  [ -f "$PALACE/$db" ] || { log "skip $db (absent)"; continue; }
  log "sqlite-backup $db"
  python3 - "$PALACE/$db" "$TMP/$db" <<'PY'
import sqlite3, sys
src_path, dst_path = sys.argv[1], sys.argv[2]
try:
    src = sqlite3.connect(f"file:{src_path}?mode=ro", uri=True)
except sqlite3.OperationalError:
    src = sqlite3.connect(src_path)          # fall back if ro open is refused
dst = sqlite3.connect(dst_path)
with dst:
    src.backup(dst)                          # atomic, consistent snapshot
dst.close(); src.close()
PY
done

restart_if_stopped
trap - EXIT INT TERM

# ── verify the COPY, not the original ───────────────────────────────────────
log "verifying copies"
CHROMA_JSON=$(py_check "$TMP/chroma.sqlite3")
log "  $CHROMA_JSON"
printf '%s' "$CHROMA_JSON" | grep -q '"quick_check": "ok"' \
  || die "chroma.sqlite3 copy failed quick_check — backup NOT committed"
KG_JSON='{}'
if [ -f "$TMP/knowledge_graph.sqlite3" ]; then
  KG_JSON=$(py_check "$TMP/knowledge_graph.sqlite3")
  log "  $KG_JSON"
  printf '%s' "$KG_JSON" | grep -q '"quick_check": "ok"' \
    || die "knowledge_graph.sqlite3 copy failed quick_check — backup NOT committed"
fi

# ── manifest ────────────────────────────────────────────────────────────────
MEMPALACE_VER=$("$HOME/.local/bin/mempalace" --version 2>/dev/null | head -1 || echo unknown)
python3 - "$TMP/MANIFEST.json" "$STAMP" "$MODE" "$PALACE" "$MEMPALACE_VER" \
         "$CHROMA_JSON" "$KG_JSON" <<'PY'
import json, os, socket, subprocess, sys
out, stamp, mode, palace, ver, chroma, kg = sys.argv[1:8]
def size(p):
    t = 0
    for root, _dirs, files in os.walk(p):
        for f in files:
            try: t += os.path.getsize(os.path.join(root, f))
            except OSError: pass
    return t
man = {
    "created_utc": stamp,
    "mode": mode,
    "host": socket.gethostname(),
    "palace_path": palace,
    "mempalace_version": ver.strip(),
    "bytes_total": size(os.path.dirname(out)),
    "chroma": json.loads(chroma),
    "knowledge_graph": json.loads(kg) if kg.strip() != "{}" else None,
    "index_consistency": (
        "byte-consistent (server quiesced)" if mode == "cold" else
        "sqlite is authoritative; HNSW segments may lag a few writes — "
        "restore with `mempalace repair --mode from-sqlite` if search misbehaves"
    ),
    "restore": [
        "# a backup dir IS a palace image — restore is a copy, not a procedure",
        "systemctl --user stop mempalace-serve",
        "mv ~/.mempalace/palace ~/.mempalace/palace.broken-$(date -u +%Y%m%dT%H%M%SZ)",
        "mkdir -p ~/.mempalace/palace",
        "rsync -a --exclude MANIFEST.json <backup>/ ~/.mempalace/palace/",
        "systemctl --user start mempalace-serve",
        "# hot backups only: if search misbehaves or the index count looks stale,",
        "# rebuild the HNSW index from the authoritative sqlite -- this also",
        "# re-creates collections with modern HNSW defaults (100/1000):",
        "mempalace repair --mode from-sqlite --archive-existing --yes",
    ],
}
with open(out, "w") as fh:
    json.dump(man, fh, indent=2)
print(json.dumps({k: man[k] for k in ("created_utc", "mode", "bytes_total")}))
PY

mv "$TMP" "$OUT"
log "committed $OUT ($(du -sh "$OUT" | cut -f1))"

# ── retention: only ever prune AFTER a committed, verified backup ───────────
if [ "$KEEP" -gt 0 ]; then
  mapfile -t all < <(find "$DEST" -maxdepth 1 -type d -name 'mp-*' | sort)
  n=${#all[@]}
  if [ "$n" -gt "$KEEP" ]; then
    for old in "${all[@]:0:$((n - KEEP))}"; do
      log "pruning $old"
      rm -rf "$old"
    done
  fi
  log "retained $(find "$DEST" -maxdepth 1 -type d -name 'mp-*' | wc -l)/$KEEP"
fi

# Leftover .incomplete dirs mean a previous run died — surface, don't hide.
for stale in "$DEST"/mp-*.incomplete; do
  [ -e "$stale" ] && log "NOTE: stale incomplete backup present: $stale"
done

log "done"