#!/usr/bin/env bash
#
# Wallarm Cloud 1.5.0, Migration process
#
# The script is designed for Wallarm Cloud version 1.5.0.
#
# The script does the following:
#   1. Detects the primary PostgreSQL pod
#   2. Dumps globals (roles, tablespaces) without role passwords
#   3. Dumps each database except those in EXCLUDE_DBS (data + schema create)
#
# Run the script from inside the wctl container.
#

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

NS="wallarm-databases"                 # your namespace
STSPREFIX="postgresql-postgresql"      # StatefulSet/pod prefix
CONTAINER="postgresql"                 # container name inside the pod
DB_USER="postgres"                     # superuser to dump with (or "wallarm" if that's your superuser)
LOCAL_OUT="./pg_dump_selected_$(date +%F_%H%M%S).sql"   # local dump path
REMOTE_OUT="/tmp/pg_dump_selected.sql"                  # file path inside the pod
REMOTE_DIR="/tmp/pgdump_parts"

# Databases to exclude
EXCLUDE_DBS=("repmgr" "postgres" "template0" "template1")

log_message() {
    local severity="$1"
    local message="$2"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [$severity] $message"
}

print_separator() {
    echo "==================================================================================================================================="
}

print_separator
log_message INFO "Finding primary in ${NS}/${STSPREFIX}-N..."

PRIMARY_POD=""
# list pods like postgresql-postgresql-0, postgresql-postgresql-1 (skip pgpool/witness)
for pod in $(kubectl -n "$NS" get pods -o name | grep "pod/${STSPREFIX}-[0-9]\+$" | sed 's#pod/##'); do
  printf " • Checking %s ... " "$pod"
  # Primary if NOT in recovery
  if kubectl -n "$NS" exec "$pod" -c "$CONTAINER" -- \
       psql -U "$DB_USER" -tAc 'SELECT NOT pg_is_in_recovery();' | grep -q '^t'; then
    PRIMARY_POD="$pod"
    break
  fi
done

if [[ -z "$PRIMARY_POD" ]]; then
  log_message ERROR "Could not find a primary pod. Aborting."
  exit 1
fi

log_message INFO "Primary is: $PRIMARY_POD"
print_separator

# Prepare remote workspace
log_message INFO "Preparing remote workspace at ${REMOTE_DIR} ..."
kubectl -n "$NS" exec "$PRIMARY_POD" -c "$CONTAINER" -- bash -lc "rm -rf '${REMOTE_DIR}' '${REMOTE_OUT}' && mkdir -p '${REMOTE_DIR}'"

# 1) Dump globals (no role passwords)
log_message INFO "Dumping globals (roles/tablespaces) without passwords ..."
kubectl -n "$NS" exec "$PRIMARY_POD" -c "$CONTAINER" -- bash -lc \
  "PGPASSWORD=\${POSTGRES_PASSWORD:-\${POSTGRES_POSTGRES_PASSWORD:-}} /opt/bitnami/postgresql/bin/pg_dumpall --globals-only --no-role-passwords -U '${DB_USER}' > '${REMOTE_OUT}'"

# Build NOT IN (...) list for SQL
EXCL_SQL=$(printf "'%s'," "${EXCLUDE_DBS[@]}")
EXCL_SQL="${EXCL_SQL%,}"

# 2) Enumerate databases to dump
log_message INFO "Enumerating databases (excluding: ${EXCLUDE_DBS[*]}) ..."
DBS=$(kubectl -n "$NS" exec "$PRIMARY_POD" -c "$CONTAINER" -- \
  psql -U "$DB_USER" -tA -c "SELECT datname FROM pg_database WHERE datistemplate = false AND datname NOT IN (${EXCL_SQL});")

if [[ -z "${DBS// }" ]]; then
  log_message WARN "No databases to dump after applying exclusions."
fi

# 3) Dump each allowed DB with --inserts and -C (CREATE DATABASE + \connect)
for db in $DBS; do
  log_message INFO "Dumping database: ${db}"
  kubectl -n "$NS" exec "$PRIMARY_POD" -c "$CONTAINER" -- bash -lc \
    "PGPASSWORD=\${POSTGRES_PASSWORD:-\${POSTGRES_POSTGRES_PASSWORD:-}} \
     /opt/bitnami/postgresql/bin/pg_dump -U '${DB_USER}' -d '${db}' --inserts -C --no-owner --no-privileges > '${REMOTE_DIR}/${db}.sql'"
done

# 4) Concatenate parts into a single file on the pod
log_message INFO "Concatenating globals + per-DB dumps into ${REMOTE_OUT} ..."
kubectl -n "$NS" exec "$PRIMARY_POD" -c "$CONTAINER" -- bash -lc \
  "for f in \$(ls -1 '${REMOTE_DIR}'/*.sql 2>/dev/null | sort); do cat \"\$f\" >> '${REMOTE_OUT}'; echo -e '\n' >> '${REMOTE_OUT}'; done"

# 5) Filter output
log_message INFO "Filtering out unwanted roles from globals (repmgr, sr_check_user)..."
kubectl -n "$NS" exec "$PRIMARY_POD" -c "$CONTAINER" -- bash -lc "
  awk '
    BEGIN {
      skip=0;
      roles=\"repmgr|wallarm|postgres|sr_check_user\";
      cre=\"^[[:space:]]*CREATE[[:space:]]+ROLE[[:space:]]+(\"roles\")\";
      alt=\"^[[:space:]]*ALTER[[:space:]]+ROLE[[:space:]]+(\"roles\")\";
    }
    {
      if (skip) {
        if (/;[[:space:]]*$/) skip=0;
        next;
      }
      if (\$0 ~ cre) { skip=1; next; }
      if (\$0 ~ alt) next;
      print;
    }
  ' '${REMOTE_OUT}' > '${REMOTE_OUT}.filtered' && mv '${REMOTE_OUT}.filtered' '${REMOTE_OUT}'
"

# 6) Copy to local
log_message INFO "Copying dump to local: ${LOCAL_OUT}"
kubectl -n "$NS" cp "$PRIMARY_POD:$REMOTE_OUT" "$LOCAL_OUT"

# 7) Cleanup
log_message INFO "Cleaning up remote temp files..."
kubectl -n "$NS" exec "$PRIMARY_POD" -c "$CONTAINER" -- bash -lc "rm -rf '${REMOTE_DIR}' '${REMOTE_OUT}'"

# 7) Renaming
log_message INFO "Renaming Databases"
sed -i -E '/^\\connect /s/_/-/g' "$LOCAL_OUT"

print_separator
log_message INFO "Done. Dump saved to: $LOCAL_OUT"