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

WORKSPACE="${WORKSPACE:-/home/openclaw/.openclaw/workspace}"
ENV_FILE="${MEMORY_RESTORE_ENV_FILE:-$WORKSPACE/.env}"

log(){
  echo "[$(date -u +'%F %T UTC')] $*"
}

if [[ -f "$ENV_FILE" ]]; then
  set -a
  # shellcheck disable=SC1090
  source "$ENV_FILE"
  set +a
else
  log "warn: env file not found: $ENV_FILE"
fi

if ! command -v psql >/dev/null 2>&1; then
  log "error: psql is not installed; cannot restore memory from Postgres"
  exit 3
fi

CONN_ARGS=()
if [[ -n "${MEMORY_DB_URL:-}" ]]; then
  CONN_ARGS+=("$MEMORY_DB_URL")
else
  missing=()
  for v in MEMORY_DB_HOST MEMORY_DB_PORT MEMORY_DB_NAME MEMORY_DB_USER MEMORY_DB_PASSWORD; do
    [[ -n "${!v:-}" ]] || missing+=("$v")
  done

  if (( ${#missing[@]} > 0 )); then
    log "warn: missing DB env vars: ${missing[*]}"
    log "warn: restore skipped (graceful fallback)"
    exit 2
  fi

  export PGHOST="$MEMORY_DB_HOST"
  export PGPORT="$MEMORY_DB_PORT"
  export PGDATABASE="$MEMORY_DB_NAME"
  export PGUSER="$MEMORY_DB_USER"
  export PGPASSWORD="$MEMORY_DB_PASSWORD"
  export PGSSLMODE="${MEMORY_DB_SSLMODE:-prefer}"
fi

mkdir -p "$WORKSPACE/memory/backups"
TS="$(date -u +'%Y%m%dT%H%M%SZ')"
BACKUP_DIR="$WORKSPACE/memory/backups/restore-$TS"
mkdir -p "$BACKUP_DIR"

# Backup current files before restore.
if [[ -f "$WORKSPACE/MEMORY.md" ]]; then
  cp "$WORKSPACE/MEMORY.md" "$BACKUP_DIR/MEMORY.md"
fi

if [[ -d "$WORKSPACE/memory" ]]; then
  find "$WORKSPACE/memory" -maxdepth 1 -type f -name '*.md' -print0 | while IFS= read -r -d '' f; do
    cp "$f" "$BACKUP_DIR/$(basename "$f")"
  done
fi

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

SQL="SELECT source_path, encode(convert_to(content, 'UTF8'), 'base64') AS content_b64 FROM mem_items WHERE source_path = 'MEMORY.md' OR source_path ~ '^memory/[^/]+\\.md$' ORDER BY source_path;"

if ! psql "${CONN_ARGS[@]}" -X -A -t -F $'\t' -c "$SQL" > "$TMP_OUT"; then
  log "error: Postgres query failed"
  exit 4
fi

if [[ ! -s "$TMP_OUT" ]]; then
  log "warn: no rows returned from mem_items for memory files"
  exit 5
fi

python3 - <<'PY' "$WORKSPACE" "$TMP_OUT"
import base64
import pathlib
import sys

workspace = pathlib.Path(sys.argv[1]).resolve()
rows_file = pathlib.Path(sys.argv[2])
restored = 0

for raw in rows_file.read_text(encoding='utf-8', errors='replace').splitlines():
    if not raw.strip():
        continue
    try:
        source_path, b64 = raw.split('\t', 1)
    except ValueError:
        continue

    source_path = source_path.strip()

    allowed = source_path == 'MEMORY.md' or (
        source_path.startswith('memory/') and source_path.endswith('.md') and '/' not in source_path[len('memory/'):]
    )
    if not allowed:
        continue

    target = (workspace / source_path).resolve()
    if workspace not in target.parents and target != workspace / 'MEMORY.md':
        continue

    target.parent.mkdir(parents=True, exist_ok=True)
    content = base64.b64decode(b64.encode('ascii'))
    target.write_bytes(content)
    restored += 1

print(f"restored_files={restored}")
if restored == 0:
    raise SystemExit(6)
PY

log "ok: memory restore completed (backup: $BACKUP_DIR)"