#!/bin/bash

# === CONFIG ===
REMOTE_USER="doc"
REMOTE_PGDATA="/var/lib/postgresql/16/main"
REMOTE_BASE="/var/lib/postgresql/16/base_restore"
REMOTE_WAL="/var/lib/postgresql/16/wal_archive"
BACKUP_SERVER="backup.sshjunkie.com"
BACKUP_BASE="/mnt/backup/pgdumps/"
BACKUP_WAL="/mnt/backups/wal/"
TIMESTAMP_FILE="$HOME/genesis-tools/miscellaneous/pitr_reference_time.txt"
LOG_FILE="$HOME/pitr_logs/pitr_staging.log"
TIMESTAMP_NOW=$(date "+%Y-%m-%d %H:%M:%S")

mkdir -p "$(dirname "$LOG_FILE")"

# === Choose Target Node ===
if [[ "$1" == "--target" && -n "$2" ]]; then
  if [[ "$2" == "db3" ]]; then
    REMOTE_HOST="replica.db3.sshjunkie.com"
  elif [[ "$2" == "db4" ]]; then
    REMOTE_HOST="replica.db4.sshjunkie.com"
  else
    echo "❌ Invalid target. Use --target db3 or --target db4."
    exit 1
  fi
else
  echo "Usage: $0 --target [db3|db4]"
  exit 1
fi

# === Validate timestamp file exists ===
if [[ ! -f "$TIMESTAMP_FILE" ]]; then
  echo "❌ No reference timestamp found at $TIMESTAMP_FILE"
  exit 1
fi

RECOVERY_TIME=$(awk '{print $1 " " $2}' "$TIMESTAMP_FILE")

echo "[$TIMESTAMP_NOW] === PITR Restore to Staging Node: $REMOTE_HOST ===" | tee -a "$LOG_FILE"

# === Find Latest Base Backup Folder on Backup Server ===
LATEST_BACKUP_FOLDER=$(ssh "$BACKUP_SERVER" "ls -td $BACKUP_BASE*/ | head -n1" | xargs basename)

if [[ -z "$LATEST_BACKUP_FOLDER" ]]; then
  echo "❌ No base backups found on backup server!"
  exit 1
fi

echo "[*] Using latest base backup folder: $LATEST_BACKUP_FOLDER"

# === Check freshness of latest backup ===
LATEST_BACKUP_TIME=$(ssh "$BACKUP_SERVER" "stat -c %Y $BACKUP_BASE$LATEST_BACKUP_FOLDER")
CURRENT_TIME=$(date +%s)
AGE_HOURS=$(( (CURRENT_TIME - LATEST_BACKUP_TIME) / 3600 ))

if [[ "$AGE_HOURS" -gt 24 ]]; then
  echo "⚠️ WARNING: Latest base backup is older than 24 hours! ($AGE_HOURS hours old)"
fi

# === Rsync Base Backup and WALs from Backup Server Directly to Replica ===
echo "[*] Rsyncing base backup directly to $REMOTE_HOST..." | tee -a "$LOG_FILE"
rsync -avz --delete "$BACKUP_SERVER:$BACKUP_BASE$LATEST_BACKUP_FOLDER/" "$REMOTE_USER@$REMOTE_HOST:$REMOTE_BASE/" >> "$LOG_FILE" 2>&1

echo "[*] Rsyncing WALs directly to $REMOTE_HOST..." | tee -a "$LOG_FILE"
rsync -avz --delete "$BACKUP_SERVER:$BACKUP_WAL" "$REMOTE_USER@$REMOTE_HOST:$REMOTE_WAL/" >> "$LOG_FILE" 2>&1

# === SSH into Replica Node and Perform Restore ===
ssh "$REMOTE_USER@$REMOTE_HOST" bash << EOF
set -e

echo "[*] Stopping PostgreSQL service..."
sudo -n systemctl stop postgresql@16-main

echo "[*] Cleaning old PGDATA directory..."
sudo -n rm -rf $REMOTE_PGDATA/*
sudo -n mkdir -p $REMOTE_PGDATA
sudo -n chown postgres:postgres $REMOTE_PGDATA

echo "[*] Copying base backup into PGDATA..."
sudo -n cp -a $REMOTE_BASE/* $REMOTE_PGDATA/
sudo -n chown -R postgres:postgres $REMOTE_PGDATA

echo "[*] Cleaning up leftover recovery artifacts..."
sudo -n rm -f $REMOTE_PGDATA/postmaster.pid $REMOTE_PGDATA/standby.signal $REMOTE_PGDATA/recovery.signal

echo "[*] Writing recovery configuration..."
sudo -n bash -c "cat > $REMOTE_PGDATA/postgresql.auto.conf << EOF2
restore_command = 'cp $REMOTE_WAL/%f %p'
recovery_target_time = '$RECOVERY_TIME'
EOF2"
sudo -n touch $REMOTE_PGDATA/recovery.signal

echo "[*] Starting PostgreSQL..."
sudo -n systemctl start postgresql@16-main

sleep 5

echo "[*] Checking recovery status..."
RECOVERY_STATE=\$(sudo -n -u postgres psql -U postgres -tAc "SELECT pg_is_in_recovery();")

if [[ "\$RECOVERY_STATE" == "f" ]]; then
  echo "[✓] Database exited recovery mode successfully."
else
  echo "[!] Still in recovery. Promoting..."
  sudo -n systemctl restart postgresql@16-main
  sleep 5
fi

echo "[*] Final replay point after recovery:"
sudo -n -u postgres psql -U postgres -tAc "SELECT pg_last_wal_replay_lsn(), now();"

EOF

echo "[✓] PITR validation run complete for $REMOTE_HOST." | tee -a "$LOG_FILE"
