#!/usr/bin/env bash
# Packet Station restore. Runs on the docker host. Restores a ps-backup.sh dump into a fresh
# database on a stack's postgres container, optionally applies the globals file, copies assets
# back into an assets volume, and prints the row count of every table at the end.
#
# It never stops, restarts or reconfigures a stack. Pointing the app at the restored database
# (or renaming it into place) is a separate, deliberate step: see scripts/backup/README.md.
#
# Usage:
#   ps-restore.sh --container NAME | --project PROJECT
#                 [--dump FILE --db NAME [--replace]]
#                 [--globals FILE]
#                 [--assets DIR_OR_TAR --assets-volume VOLUME]
#                 [--verify-assets --db NAME --assets-volume VOLUME]
#                 [--counts-only --db NAME]
#
#   --container / --project  the postgres container, by name or by compose project.
#   --dump FILE              a ps-*.dump file (pg_dump -Fc).
#   --db NAME                the database to restore into. It must not exist, unless --replace,
#                            which drops it first. DROP DATABASE refuses while anyone is connected.
#   --owner ROLE             owner of the new database (default ps_migrate, as initdb makes it).
#   --app-role ROLE          role given CONNECT and TEMPORARY on it (default ps_app).
#   --globals FILE           a ps-*.globals.sql file. Roles that already exist are reported and
#                            left alone. The file holds no passwords: set them from the secrets.
#   --assets SRC             a folder, or a .tar / .tar.gz / .tgz, holding the files that belong
#                            under ASSETS_DIR. Copied into --assets-volume through a throwaway
#                            container (no network) that runs the postgres image already on the host.
#                            Existing files are never overwritten.
#   --assets-volume VOLUME   the docker volume to copy into (for example packetstation-test_assets).
#   --assets-owner UID:GID   owner for the copied files (default 10001:10001, the image's ps user).
#   --verify-assets          compare sha256 of every file in --assets-volume with assets.sha256 in --db.
#   --counts-only            only print the row counts of --db.
#
# Connects inside the container as the postgres superuser over the local socket. No password
# is read, passed or printed.

set -euo pipefail
umask 077

PG_USER="${PG_USER:-postgres}"
DOCKER="${DOCKER:-docker}"

log() { printf '%s ps-restore: %s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)" "$*" >&2; }
die() { log "$*"; exit 1; }
usage() { sed -n '2,36p' "$0" | sed 's/^# \{0,1\}//' >&2; exit 2; }

container=""
project=""
dump=""
db=""
replace=0
owner="ps_migrate"
app_role="ps_app"
globals=""
assets=""
assets_volume=""
assets_owner="10001:10001"
verify_assets=0
counts_only=0

while [[ $# -gt 0 ]]; do
  case "$1" in
    --container) container="${2:?}"; shift 2 ;;
    --project) project="${2:?}"; shift 2 ;;
    --dump) dump="${2:?}"; shift 2 ;;
    --db) db="${2:?}"; shift 2 ;;
    --replace) replace=1; shift ;;
    --owner) owner="${2:?}"; shift 2 ;;
    --app-role) app_role="${2:?}"; shift 2 ;;
    --globals) globals="${2:?}"; shift 2 ;;
    --assets) assets="${2:?}"; shift 2 ;;
    --assets-volume) assets_volume="${2:?}"; shift 2 ;;
    --assets-owner) assets_owner="${2:?}"; shift 2 ;;
    --verify-assets) verify_assets=1; shift ;;
    --counts-only) counts_only=1; shift ;;
    -h|--help) usage ;;
    *) log "unknown argument $1"; usage ;;
  esac
done

ident_re='^[a-z_][a-z0-9_]{0,62}$'
[[ -z "$db" || "$db" =~ $ident_re ]] || die "--db must be a plain lower-case name (letters, digits, _)"
[[ "$owner" =~ $ident_re ]] || die "--owner must be a plain lower-case role name"
[[ "$app_role" =~ $ident_re ]] || die "--app-role must be a plain lower-case role name"
[[ "$assets_owner" =~ ^[0-9]+:[0-9]+$ ]] || die "--assets-owner must be UID:GID"

if [[ -n "$project" ]]; then
  [[ -z "$container" ]] || die "use --container or --project, not both"
  container="$("$DOCKER" ps -q \
    --filter "label=com.docker.compose.project=$project" \
    --filter "label=com.docker.compose.service=postgres")"
  [[ -n "$container" ]] || die "no running postgres container for compose project $project"
  [[ "$(printf '%s\n' "$container" | wc -l)" -eq 1 ]] || die "more than one postgres container for $project"
fi
[[ -n "$container" ]] || die "--container or --project is required"

if [[ -n "$dump" || $counts_only -eq 1 || $verify_assets -eq 1 ]]; then
  [[ -n "$db" ]] || die "--db is required"
fi
if [[ -n "$assets" || $verify_assets -eq 1 ]]; then
  [[ -n "$assets_volume" ]] || die "--assets-volume is required"
fi
if [[ -z "$dump" && -z "$globals" && -z "$assets" && $verify_assets -eq 0 && $counts_only -eq 0 ]]; then
  die "nothing to do: give --dump, --globals, --assets, --verify-assets or --counts-only"
fi

# psql inside the container, reading SQL from stdin. Variables go through -v and :'name'.
psql_c() {
  local dbname="$1"; shift
  "$DOCKER" exec -i "$container" psql -X -q -v ON_ERROR_STOP=1 -U "$PG_USER" -d "$dbname" "$@"
}

db_exists() {
  local n
  n="$(psql_c postgres -At -v name="$1" <<'SQL'
SELECT count(*) FROM pg_database WHERE datname = :'name';
SQL
)"
  [[ "$n" == "1" ]]
}

role_exists() {
  local n
  n="$(psql_c postgres -At -v name="$1" <<'SQL'
SELECT count(*) FROM pg_roles WHERE rolname = :'name';
SQL
)"
  [[ "$n" == "1" ]]
}

print_counts() {
  local dbname="$1"
  log "row counts in $dbname"
  psql_c "$dbname" -A -F $'\t' -P footer=off <<'SQL'
SELECT c.relname AS table_name,
       (xpath('/row/n/text()',
              query_to_xml(format('SELECT count(*) AS n FROM %I.%I', n.nspname, c.relname), false, true, '')
       ))[1]::text::bigint AS row_count
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = 'public' AND c.relkind IN ('r', 'p')
ORDER BY c.relname;
SQL
}

# A throwaway container on no network, running the image the postgres container already uses,
# so nothing is downloaded.
helper_image() {
  "$DOCKER" inspect -f '{{.Config.Image}}' "$container"
}

# ------------------------------------------------------------------ globals

if [[ -n "$globals" ]]; then
  [[ -r "$globals" ]] || die "cannot read $globals"
  if grep -qiE "PASSWORD[[:space:]]+'(SCRAM|md5)" "$globals"; then
    die "$globals holds password hashes; ps-backup.sh writes globals without them. Refusing."
  fi
  log "applying globals from $(basename -- "$globals") (existing roles are left as they are)"
  # Without ON_ERROR_STOP: CREATE ROLE for a role that exists fails and the rest goes on.
  out="$("$DOCKER" exec -i "$container" psql -X -q -U "$PG_USER" -d postgres < "$globals" 2>&1 || true)"
  exists="$(printf '%s\n' "$out" | grep -c 'already exists' || true)"
  other="$(printf '%s\n' "$out" | grep 'ERROR' | grep -vc 'already exists' || true)"
  log "globals: $exists statements skipped (already exists), $other other errors"
  if [[ "$other" -gt 0 ]]; then
    printf '%s\n' "$out" | grep 'ERROR' | grep -v 'already exists' >&2
    die "globals had errors"
  fi
fi

# ------------------------------------------------------------------ database

if [[ -n "$dump" ]]; then
  [[ -r "$dump" ]] || die "cannot read $dump"
  started=$SECONDS
  if db_exists "$db"; then
    [[ $replace -eq 1 ]] || die "database $db exists. Give --replace to drop it first."
    log "dropping database $db (--replace)"
    psql_c postgres -v name="$db" <<'SQL'
SELECT format('DROP DATABASE %I', :'name') \gexec
SQL
  fi
  role_exists "$owner" || die "role $owner does not exist (apply --globals or start the stack once)"

  log "creating database $db owned by $owner"
  # The same shape as deploy/postgres/initdb/01-roles.sh.
  psql_c postgres -v name="$db" -v owner="$owner" <<'SQL'
SELECT format('CREATE DATABASE %I OWNER %I ENCODING %L TEMPLATE template0', :'name', :'owner', 'UTF8') \gexec
SELECT format('REVOKE ALL ON DATABASE %I FROM PUBLIC', :'name') \gexec
SQL
  if role_exists "$app_role"; then
    psql_c postgres -v name="$db" -v app="$app_role" <<'SQL'
SELECT format('GRANT CONNECT, TEMPORARY ON DATABASE %I TO %I', :'name', :'app') \gexec
SQL
  else
    log "role $app_role does not exist; no CONNECT grant made"
  fi

  log "restoring $(basename -- "$dump") ($(stat -c %s "$dump") bytes) into $db"
  # One transaction: a restore that fails leaves an empty database, never a half-restored one.
  "$DOCKER" exec -i "$container" pg_restore -U "$PG_USER" -d "$db" --exit-on-error --single-transaction \
    < "$dump" || die "pg_restore failed; database $db is left empty. Drop it with --replace on the next try."
  log "database restored in $(( SECONDS - started )) s"
fi

# ------------------------------------------------------------------ assets

if [[ -n "$assets" ]]; then
  [[ -e "$assets" ]] || die "$assets does not exist"
  "$DOCKER" volume inspect "$assets_volume" >/dev/null 2>&1 || die "volume $assets_volume does not exist"
  image="$(helper_image)"
  started=$SECONDS
  src="$(cd -- "$(dirname -- "$assets")" && pwd)/$(basename -- "$assets")"
  if [[ -d "$assets" ]]; then
    log "copying assets from folder $src into volume $assets_volume"
    "$DOCKER" run --rm --network none --user 0:0 --entrypoint sh \
      -v "$assets_volume:/dst" -v "$src:/src:ro" "$image" -c \
      'set -eu; cp -a -n /src/. /dst/; chown -R "$1" /dst' sh "$assets_owner"
  else
    case "$assets" in
      *.tar.gz|*.tgz) tar_flags="-xzf" ;;
      *.tar) tar_flags="-xf" ;;
      *) die "--assets must be a folder, .tar, .tar.gz or .tgz" ;;
    esac
    log "copying assets from $(basename -- "$src") into volume $assets_volume"
    "$DOCKER" run --rm --network none --user 0:0 --entrypoint sh \
      -v "$assets_volume:/dst" -v "$src:/src.tar:ro" "$image" -c \
      'set -eu; tar "$1" /src.tar -C /dst --skip-old-files --no-same-owner; chown -R "$2" /dst' \
      sh "$tar_flags" "$assets_owner"
  fi
  log "assets copied in $(( SECONDS - started )) s"
fi

# ------------------------------------------------------------------ verify assets

if [[ $verify_assets -eq 1 ]]; then
  "$DOCKER" volume inspect "$assets_volume" >/dev/null 2>&1 || die "volume $assets_volume does not exist"
  image="$(helper_image)"
  work="$(mktemp -d)"
  trap 'rm -rf -- "$work"' EXIT
  # sha256sum check format: "<hex>  <path>".
  psql_c "$db" -At <<'SQL' > "$work/assets.sha256"
SELECT encode(sha256, 'hex') || '  ' || storage_key FROM assets ORDER BY storage_key;
SQL
  rows="$(wc -l < "$work/assets.sha256")"
  chmod 0755 "$work"
  chmod 0644 "$work/assets.sha256"
  log "checking $rows asset rows of $db against volume $assets_volume"
  set +e
  "$DOCKER" run --rm --network none --user 0:0 --entrypoint sh \
    -v "$assets_volume:/dst:ro" -v "$work:/work:ro" "$image" -c '
      cd /dst || exit 3
      sha256sum -c /work/assets.sha256 2>/dev/null
      find . -type f ! -path "./tmp/*" | sed "s|^\./||" | sort > /tmp/files
      cut -c67- /work/assets.sha256 | sort > /tmp/rows
      comm -23 /tmp/files /tmp/rows | sed "s|^|ORPHAN: |"
      exit 0' > "$work/result"
  rc=$?
  set -e
  [[ $rc -eq 0 ]] || die "the check container failed (exit $rc)"
  ok="$(grep -c ': OK$' "$work/result" || true)"
  bad="$(grep -c ': FAILED$' "$work/result" || true)"
  missing="$(grep -c ': FAILED open or read$' "$work/result" || true)"
  orphans="$(grep -c '^ORPHAN: ' "$work/result" || true)"
  log "assets: $rows rows, $ok match, $bad differ, $missing missing, $orphans files with no row"
  grep -v ': OK$' "$work/result" >&2 || true
  if [[ "$ok" -ne "$rows" ]]; then
    die "asset check failed"
  fi
fi

# ------------------------------------------------------------------ counts

if [[ -n "$db" && ( -n "$dump" || $counts_only -eq 1 ) ]]; then
  print_counts "$db"
fi