#!/usr/bin/env bash
#
# Wallarm Cloud 1.5.0, Migration restore
#
# The script:
#   1. Detects the PRIMARY pod using CNPG Cluster status
#   2. Copies a local SQL dump into the PRIMARY pod
#   3. Restores the dump with psql (ON_ERROR_STOP)
#
# Run the script from inside the wctl container, or anywhere with kubectl access.
#

set -euo pipefail

# --- Config -------------------------------------------------------------------
NS="wallarm-databases"                       # namespace
CLUSTER_NAME="postgresql"                    # CNPG Cluster name (kubectl get cluster -n ...)
LOCAL_DUMP_PATH="${LOCAL_DUMP_PATH:-./pg_dump_selected_latest.sql}"
REMOTE_DUMP_PATH="/var/lib/postgresql/data/restore.sql"          # path inside the pod
PSQL="/opt/bitnami/postgresql/bin/psql"      # Bitnami psql path
# ------------------------------------------------------------------------------

log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] [$1] $2"; }
sep() { echo "==================================================================================================================================="; }

require_file() {
  if [[ ! -s "$1" ]]; then
    log ERROR "Dump file '$1' not found or empty"; exit 1
  fi
}

get_primary_from_json() {
  # CNPG usually exposes .status.currentPrimary
  kubectl -n "$NS" get cluster "$CLUSTER_NAME" -o jsonpath='{.status.currentPrimary}' 2>/dev/null || true
}

get_primary_from_table() {
  # Fallback: parse the PRIMARY column from the wide table output
  kubectl -n "$NS" get cluster | awk -v c="$CLUSTER_NAME" '$1==c {print $NF}'
}

# --- Preflight ----------------------------------------------------------------
sep
log INFO "Validating local dump file..."
require_file "$LOCAL_DUMP_PATH"

sep
log INFO "Detecting PRIMARY pod for Cluster '$CLUSTER_NAME' in namespace '$NS'..."

PRIMARY_POD="$(get_primary_from_json)"
if [[ -z "${PRIMARY_POD}" ]]; then
  log WARN "Could not get primary via jsonpath; parsing table output..."
  PRIMARY_POD="$(get_primary_from_table)"
fi

if [[ -z "${PRIMARY_POD}" ]]; then
  log ERROR "Unable to detect PRIMARY pod. Aborting."
  exit 1
fi

log INFO "PRIMARY (CNPG) reports: ${PRIMARY_POD}"

# CNPG shows pod names like 'postgresql-1'. Your actual K8s pod may be 'postgresql-postgresql-1'.
# If your StatefulSet uses a prefix, map it here:
# Try a few common patterns automatically; pick the first that exists.
CANDIDATE_PODS=(
  "${PRIMARY_POD}"
  "postgresql-postgresql-${PRIMARY_POD##*-}"   # map 'postgresql-1' -> 'postgresql-postgresql-1'
  "postgresql-${PRIMARY_POD##*-}"              # map 'postgresql-1' -> 'postgresql-1' (as-is)
)

REAL_POD=""
for p in "${CANDIDATE_PODS[@]}"; do
  if kubectl -n "$NS" get pod "$p" >/dev/null 2>&1; then
    REAL_POD="$p"
    break
  fi
done

if [[ -z "$REAL_POD" ]]; then
  log ERROR "Could not resolve actual Pod name from PRIMARY='${PRIMARY_POD}'. Checked: ${CANDIDATE_PODS[*]}"
  exit 1
fi

log INFO "Using Pod: ${REAL_POD}"
sep

# --- Copy dump into the pod ---------------------------------------------------
log INFO "Copying local dump to pod: ${LOCAL_DUMP_PATH} -> ${REAL_POD}:${REMOTE_DUMP_PATH}"
kubectl -n "$NS" cp "$LOCAL_DUMP_PATH" "${REAL_POD}:${REMOTE_DUMP_PATH}"


# --- Extract USER && PASSWORD FROM DATABASES ----------------------------------
log INFO "Copying local dump to pod: ${LOCAL_DUMP_PATH} -> ${REAL_POD}:${REMOTE_DUMP_PATH}"
export POSTGRES_USER=$(kubectl -n wallarm-databases get secret postgres-wallarm -o jsonpath='{.data.username}' | base64 --decode)
export POSTGRES_PASSWORD=$(kubectl -n wallarm-databases get secret postgres-wallarm -o jsonpath='{.data.password}' | base64 --decode)


# --- Restore ------------------------------------------------------------------
log INFO "Starting restore with psql (ON_ERROR_STOP=1)..."
log INFO "Tips: consider pausing writers, and ensure this PRIMARY is in RW mode."

# Speed/robust settings applied per-session.
# - ON_ERROR_STOP stops on first error
# - client_min_messages to notice
# - synchronous_commit OFF can speed up big restores (optional, uncomment if desired)
# - maintenance_work_mem/temp_buffers can be tuned if needed

kubectl -n "$NS" exec "$REAL_POD" -- bash -lc "
  set -euo pipefail
  PGPASSWORD="${POSTGRES_PASSWORD}" ${PSQL} -U '${POSTGRES_USER}' -h 127.0.0.1  \
    -v client_min_messages=notice \
    -d postgres \
    -f '${REMOTE_DUMP_PATH}'
"

log INFO "Restore finished successfully."
sep

# --- Optional cleanup ---------------------------------------------------------
kubectl -n "$NS" exec "$REAL_POD" -- rm -f "$REMOTE_DUMP_PATH" || true
log INFO "Removed ${REMOTE_DUMP_PATH} from pod."

log INFO "All done."