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

# Backup + WAL archiving setup for Patroni/PostgreSQL DB nodes.
#
# What this sets up:
# - pgBackRest repository on a backup host
# - pgBackRest client config on db1/db2/db3
# - WAL archiving from PostgreSQL to the repo
# - daily full backup + frequent incremental backup via cron
#
# Usage:
#   sudo ./setup-backup.sh /path/to/backup.env
#
# ROLE values:
#   ROLE=repo   -> backup repository server
#   ROLE=db     -> database node
#
# Assumptions:
# - Debian/Ubuntu
# - db1/db2/db3/backup1 resolve via /etc/hosts or DNS
# - SSH connectivity works between db nodes and repo host over WireGuard
# - PostgreSQL 16 binaries live in /usr/lib/postgresql/16/bin
# - Patroni data dir is /var/lib/postgresql/data
#
# This script:
# - installs pgBackRest
# - writes pgBackRest config
# - creates systemd timer or cron entries for backups on repo host
# - writes Patroni-compatible PostgreSQL archive settings on DB hosts
#
# IMPORTANT:
# - After applying DB config changes, restart Patroni on each DB node one at a time.
# - Run stanza-create from the repo host after all DB configs are present.

ENV_FILE="${1:-}"
if [[ -z "$ENV_FILE" || ! -f "$ENV_FILE" ]]; then
  echo "Usage: sudo $0 /path/to/backup.env"
  exit 1
fi

set -a
source "$ENV_FILE"
set +a

require() {
  local var="$1"
  if [[ -z "${!var:-}" ]]; then
    echo "Missing required variable: $var" >&2
    exit 1
  fi
}

ROLE="${ROLE:-}"
if [[ "$ROLE" != "repo" && "$ROLE" != "db" ]]; then
  echo "ROLE must be repo or db" >&2
  exit 1
fi

common_vars=(
  STANZA REPO_HOST REPO_PATH REPO_RETENTION_FULL REPO_RETENTION_DIFF
  DB1_HOST DB2_HOST DB3_HOST
  PG_SUPERUSER PG_SUPERUSER_PASSWORD
)
for v in "${common_vars[@]}"; do require "$v"; done

if [[ "$ROLE" == "repo" ]]; then
  repo_vars=(REPO_HOSTNAME)
  for v in "${repo_vars[@]}"; do require "$v"; done
else
  db_vars=(NODE_NAME PATRONI_CONFIG_PATH PGBACKREST_DB_PATH)
  for v in "${db_vars[@]}"; do require "$v"; done
fi

log() { echo "[+] $*"; }

install_packages() {
  export DEBIAN_FRONTEND=noninteractive
  apt-get update
  apt-get install -y pgbackrest cron openssh-client
}

setup_repo_host() {
  log "Setting up pgBackRest repository host"
  install -d -m 0750 /etc/pgbackrest
  install -d -m 0750 "${REPO_PATH}"
  chown -R postgres:postgres "${REPO_PATH}" || true

  cat > /etc/pgbackrest/pgbackrest.conf <<EOF
[global]
repo1-path=${REPO_PATH}
repo1-retention-full=${REPO_RETENTION_FULL}
repo1-retention-diff=${REPO_RETENTION_DIFF}
start-fast=y
process-max=4
log-level-console=info
log-level-file=detail

[${STANZA}]
pg1-host=${DB1_HOST}
pg1-path=/var/lib/postgresql/data
pg1-port=5432
pg1-user=postgres

pg2-host=${DB2_HOST}
pg2-path=/var/lib/postgresql/data
pg2-port=5432
pg2-user=postgres

pg3-host=${DB3_HOST}
pg3-path=/var/lib/postgresql/data
pg3-port=5432
pg3-user=postgres
EOF

  chmod 640 /etc/pgbackrest/pgbackrest.conf

  cat > /usr/local/bin/pgbackrest-full.sh <<EOF
#!/usr/bin/env bash
set -euo pipefail
pgbackrest --stanza=${STANZA} backup --type=full
EOF
  chmod +x /usr/local/bin/pgbackrest-full.sh

  cat > /usr/local/bin/pgbackrest-incr.sh <<EOF
#!/usr/bin/env bash
set -euo pipefail
pgbackrest --stanza=${STANZA} backup --type=incr
EOF
  chmod +x /usr/local/bin/pgbackrest-incr.sh

  cat > /usr/local/bin/pgbackrest-check.sh <<EOF
#!/usr/bin/env bash
set -euo pipefail
pgbackrest --stanza=${STANZA} check
EOF
  chmod +x /usr/local/bin/pgbackrest-check.sh

  cat > /etc/cron.d/pgbackrest <<'EOF'
SHELL=/bin/bash
PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin

# Daily full backup at 02:15
15 2 * * * root /usr/local/bin/pgbackrest-full.sh >> /var/log/pgbackrest-full.log 2>&1

# Incremental backup every 6 hours
10 */6 * * * root /usr/local/bin/pgbackrest-incr.sh >> /var/log/pgbackrest-incr.log 2>&1

# Daily check at 01:40
40 1 * * * root /usr/local/bin/pgbackrest-check.sh >> /var/log/pgbackrest-check.log 2>&1
EOF

  systemctl enable cron
  systemctl restart cron

  cat <<EOF

Repository host setup complete.

Next steps:
1. Make sure postgres SSH access from this host to db1/db2/db3 works:
   sudo -u postgres ssh ${DB1_HOST} true
   sudo -u postgres ssh ${DB2_HOST} true
   sudo -u postgres ssh ${DB3_HOST} true

2. After DB nodes are configured, create the stanza:
   sudo -u postgres pgbackrest --stanza=${STANZA} stanza-create

3. Run an initial full backup:
   sudo -u postgres pgbackrest --stanza=${STANZA} backup --type=full
EOF
}

setup_db_host() {
  log "Setting up pgBackRest client on DB node"
  install -d -m 0750 /etc/pgbackrest

  cat > /etc/pgbackrest/pgbackrest.conf <<EOF
[global]
repo1-host=${REPO_HOST}
repo1-host-user=postgres
repo1-path=${REPO_PATH}
log-level-console=info
log-level-file=detail
process-max=2

[${STANZA}]
pg1-path=${PGBACKREST_DB_PATH}
EOF

  chmod 640 /etc/pgbackrest/pgbackrest.conf

  install -d -m 0700 /var/lib/postgresql/.ssh || true
  chown -R postgres:postgres /var/lib/postgresql/.ssh || true

  cat > /usr/local/bin/apply-patroni-backup-settings.sh <<EOF
#!/usr/bin/env bash
set -euo pipefail

CFG="${PATRONI_CONFIG_PATH}"

python3 - <<'PY'
from pathlib import Path
cfg = Path("${PATRONI_CONFIG_PATH}")
text = cfg.read_text()

required = [
    'archive_mode: "on"',
    'archive_command: \'pgbackrest --stanza=${STANZA} archive-push %p\''
]

if 'archive_mode:' not in text:
    marker = '        wal_keep_size: 256MB\\n'
    if marker in text:
        text = text.replace(marker, marker + '        archive_mode: "on"\\n        archive_command: \\'pgbackrest --stanza=${STANZA} archive-push %p\\'\\n')
    else:
        raise SystemExit('Could not find insertion point for archive settings')

if 'restore_command:' not in text:
    marker = '  parameters:\\n    unix_socket_directories: /var/run/postgresql\\n'
    repl = "  parameters:\\n    unix_socket_directories: /var/run/postgresql\\n    restore_command: 'pgbackrest --stanza=${STANZA} archive-get %f %p'\\n"
    if marker in text:
        text = text.replace(marker, repl)
    else:
        raise SystemExit('Could not find insertion point for restore_command')

cfg.write_text(text)
print("Updated Patroni config with archive settings")
PY
EOF
  chmod +x /usr/local/bin/apply-patroni-backup-settings.sh

  cat <<EOF

DB host client setup complete.

Next steps on this DB node:
1. Ensure postgres user SSH key exists:
   sudo -u postgres ssh-keygen -t ed25519 -N '' -f /var/lib/postgresql/.ssh/id_ed25519

2. Copy the public key to the repo host:
   sudo -u postgres ssh-copy-id postgres@${REPO_HOST}

3. Apply backup archive settings into Patroni YAML:
   sudo /usr/local/bin/apply-patroni-backup-settings.sh

4. Restart Patroni on this node during a maintenance window:
   cd /opt/db-stack && docker compose restart patroni
EOF
}

install_packages
if [[ "$ROLE" == "repo" ]]; then
  setup_repo_host
else
  setup_db_host
fi