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

# One-command installer for either a DB node or the proxy node.
# Usage:
#   sudo ./install-node.sh /path/to/node.env
#
# Supported ROLE values:
#   ROLE=db
#   ROLE=proxy
#
# Assumptions:
# - Debian/Ubuntu host
# - You will define db1, db2, db3, and proxy1 in /etc/hosts
# - WireGuard private keys and peer public keys are already generated
# - Run as root or with sudo

ENV_FILE="${1:-}"
if [[ -z "$ENV_FILE" || ! -f "$ENV_FILE" ]]; then
  echo "Usage: sudo $0 /path/to/node.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
}

common_required=(
  ROLE HOSTNAME_FQDN WG_ADDRESS WG_PORT WG_PRIVATE_KEY
  DB1_PUBLIC_IP DB2_PUBLIC_IP DB3_PUBLIC_IP PROXY1_PUBLIC_IP
  DB1_WG_IP DB2_WG_IP DB3_WG_IP PROXY1_WG_IP
  DB1_PUBKEY DB2_PUBKEY DB3_PUBKEY PROXY1_PUBKEY
)
for v in "${common_required[@]}"; do require "$v"; done

if [[ "$ROLE" == "db" ]]; then
  db_required=(
    NODE_NAME PATRONI_SCOPE PATRONI_NAMESPACE POSTGRES_PASSWORD
    REPLICATION_PASSWORD PATRONI_SUPERUSER_USERNAME PATRONI_REPLICATION_USERNAME
    PATRONI_DB
  )
  for v in "${db_required[@]}"; do require "$v"; done
elif [[ "$ROLE" == "proxy" ]]; then
  proxy_required=(
    PGADMIN_DEFAULT_EMAIL PGADMIN_DEFAULT_PASSWORD
    POSTGREST_DB_URI POSTGREST_DB_SCHEMAS POSTGREST_DB_ANON_ROLE
    CADDY_EMAIL PGADMIN_DOMAIN API_DOMAIN
    PGADMIN_BASIC_AUTH_USER PGADMIN_BASIC_AUTH_HASH
    API_BASIC_AUTH_USER API_BASIC_AUTH_HASH
  )
  for v in "${proxy_required[@]}"; do require "$v"; done
else
  echo "ROLE must be db or proxy" >&2
  exit 1
fi

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

install_packages() {
  export DEBIAN_FRONTEND=noninteractive
  log "Installing packages"
  apt-get update
  apt-get install -y ca-certificates curl gnupg lsb-release ufw wireguard wireguard-tools gettext-base
  if ! command -v docker >/dev/null 2>&1; then
    install -m 0755 -d /etc/apt/keyrings
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg
    chmod a+r /etc/apt/keyrings/docker.gpg
    . /etc/os-release
    echo \
      "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] \
      https://download.docker.com/linux/ubuntu ${VERSION_CODENAME} stable" \
      > /etc/apt/sources.list.d/docker.list
    apt-get update
    apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
  fi
  systemctl enable docker
  systemctl start docker
}

write_wireguard() {
  log "Writing WireGuard config"
  install -d -m 0700 /etc/wireguard

  cat > /etc/wireguard/wg0.conf <<EOF
[Interface]
Address = ${WG_ADDRESS}
ListenPort = ${WG_PORT}
PrivateKey = ${WG_PRIVATE_KEY}
SaveConfig = false

[Peer]
PublicKey = ${DB1_PUBKEY}
AllowedIPs = ${DB1_WG_IP}/32
Endpoint = ${DB1_PUBLIC_IP}:${WG_PORT}
PersistentKeepalive = 25

[Peer]
PublicKey = ${DB2_PUBKEY}
AllowedIPs = ${DB2_WG_IP}/32
Endpoint = ${DB2_PUBLIC_IP}:${WG_PORT}
PersistentKeepalive = 25

[Peer]
PublicKey = ${DB3_PUBKEY}
AllowedIPs = ${DB3_WG_IP}/32
Endpoint = ${DB3_PUBLIC_IP}:${WG_PORT}
PersistentKeepalive = 25

[Peer]
PublicKey = ${PROXY1_PUBKEY}
AllowedIPs = ${PROXY1_WG_IP}/32
Endpoint = ${PROXY1_PUBLIC_IP}:${WG_PORT}
PersistentKeepalive = 25
EOF

  chmod 600 /etc/wireguard/wg0.conf
  systemctl enable wg-quick@wg0
  systemctl restart wg-quick@wg0
}

configure_hosts() {
  log "Writing /etc/hosts entries"
  grep -qE "^[0-9.]+\s+db1(\s|$)" /etc/hosts || echo "${DB1_WG_IP} db1" >> /etc/hosts
  grep -qE "^[0-9.]+\s+db2(\s|$)" /etc/hosts || echo "${DB2_WG_IP} db2" >> /etc/hosts
  grep -qE "^[0-9.]+\s+db3(\s|$)" /etc/hosts || echo "${DB3_WG_IP} db3" >> /etc/hosts
  grep -qE "^[0-9.]+\s+proxy1(\s|$)" /etc/hosts || echo "${PROXY1_WG_IP} proxy1" >> /etc/hosts
}

configure_firewall_db() {
  log "Configuring UFW for DB node"
  ufw --force reset
  ufw default deny incoming
  ufw default allow outgoing
  ufw allow 22/tcp
  ufw allow ${WG_PORT}/udp
  ufw allow from ${DB1_WG_IP} to any port 2379 proto tcp
  ufw allow from ${DB2_WG_IP} to any port 2379 proto tcp
  ufw allow from ${DB3_WG_IP} to any port 2379 proto tcp
  ufw allow from ${DB1_WG_IP} to any port 2380 proto tcp
  ufw allow from ${DB2_WG_IP} to any port 2380 proto tcp
  ufw allow from ${DB3_WG_IP} to any port 2380 proto tcp
  ufw allow from ${DB1_WG_IP} to any port 5432 proto tcp
  ufw allow from ${DB2_WG_IP} to any port 5432 proto tcp
  ufw allow from ${DB3_WG_IP} to any port 5432 proto tcp
  ufw allow from ${PROXY1_WG_IP} to any port 5432 proto tcp
  ufw allow from ${PROXY1_WG_IP} to any port 8008 proto tcp
  ufw allow from ${DB1_WG_IP} to any port 8008 proto tcp
  ufw allow from ${DB2_WG_IP} to any port 8008 proto tcp
  ufw allow from ${DB3_WG_IP} to any port 8008 proto tcp
  ufw --force enable
}

configure_firewall_proxy() {
  log "Configuring UFW for proxy node"
  ufw --force reset
  ufw default deny incoming
  ufw default allow outgoing
  ufw allow 22/tcp
  ufw allow ${WG_PORT}/udp
  ufw allow 80/tcp
  ufw allow 443/tcp
  ufw allow 443/udp
  ufw allow 5432/tcp
  ufw allow from ${DB1_WG_IP} to any port 8080 proto tcp
  ufw allow from ${DB2_WG_IP} to any port 8080 proto tcp
  ufw allow from ${DB3_WG_IP} to any port 8080 proto tcp
  ufw allow from ${DB1_WG_IP} to any port 3000 proto tcp
  ufw allow from ${DB2_WG_IP} to any port 3000 proto tcp
  ufw allow from ${DB3_WG_IP} to any port 3000 proto tcp
  ufw --force enable
}

deploy_db_stack() {
  log "Deploying DB stack"
  mkdir -p /opt/db-stack/{postgres-data,etcd-data}
  cd /opt/db-stack

  cat > Dockerfile <<'EOF'
FROM postgres:16-bookworm
RUN apt-get update \
    && apt-get install -y --no-install-recommends python3 python3-pip python3-psycopg2 curl ca-certificates postgresql-16-cron \
    && pip3 install --no-cache-dir patroni[etcd3] \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*
COPY patroni.yml /etc/patroni/patroni.yml
CMD ["patroni", "/etc/patroni/patroni.yml"]
EOF

  cat > compose.yaml <<EOF
services:
  etcd:
    image: quay.io/coreos/etcd:v3.5.15
    container_name: ${NODE_NAME}-etcd
    restart: unless-stopped
    network_mode: host
    environment:
      ETCD_NAME: ${NODE_NAME}
      ETCD_DATA_DIR: /var/lib/etcd
      ETCD_INITIAL_ADVERTISE_PEER_URLS: http://${WG_ADDRESS%/*}:2380
      ETCD_ADVERTISE_CLIENT_URLS: http://${WG_ADDRESS%/*}:2379
      ETCD_LISTEN_PEER_URLS: http://${WG_ADDRESS%/*}:2380
      ETCD_LISTEN_CLIENT_URLS: http://${WG_ADDRESS%/*}:2379,http://127.0.0.1:2379
      ETCD_INITIAL_CLUSTER: db1=http://db1:2380,db2=http://db2:2380,db3=http://db3:2380
      ETCD_INITIAL_CLUSTER_TOKEN: pg-ha-cluster
      ETCD_INITIAL_CLUSTER_STATE: new
    volumes:
      - ./etcd-data:/var/lib/etcd

  patroni:
    build: .
    container_name: ${NODE_NAME}-patroni
    restart: unless-stopped
    network_mode: host
    depends_on:
      - etcd
    volumes:
      - ./postgres-data:/var/lib/postgresql/data
      - ./patroni.yml:/etc/patroni/patroni.yml:ro
EOF

  cat > patroni.yml <<EOF
scope: ${PATRONI_SCOPE}
namespace: ${PATRONI_NAMESPACE}
name: ${NODE_NAME}

restapi:
  listen: ${WG_ADDRESS%/*}:8008
  connect_address: ${WG_ADDRESS%/*}:8008

etcd3:
  hosts:
    - db1:2379
    - db2:2379
    - db3:2379

bootstrap:
  dcs:
    ttl: 30
    loop_wait: 10
    retry_timeout: 10
    maximum_lag_on_failover: 1048576
    postgresql:
      use_pg_rewind: true
      use_slots: true
      parameters:
        wal_level: replica
        hot_standby: "on"
        max_wal_senders: 10
        max_replication_slots: 10
        wal_keep_size: 256MB
        shared_preload_libraries: "pg_cron"
        cron.database_name: "${PATRONI_DB}"
        max_connections: 300
        shared_buffers: 256MB
  initdb:
    - encoding: UTF8
    - data-checksums

postgresql:
  listen: ${WG_ADDRESS%/*}:5432
  connect_address: ${WG_ADDRESS%/*}:5432
  data_dir: /var/lib/postgresql/data
  bin_dir: /usr/lib/postgresql/16/bin
  authentication:
    superuser:
      username: ${PATRONI_SUPERUSER_USERNAME}
      password: ${POSTGRES_PASSWORD}
    replication:
      username: ${PATRONI_REPLICATION_USERNAME}
      password: ${REPLICATION_PASSWORD}
  pg_hba:
    - host replication ${PATRONI_REPLICATION_USERNAME} ${DB1_WG_IP}/32 scram-sha-256
    - host replication ${PATRONI_REPLICATION_USERNAME} ${DB2_WG_IP}/32 scram-sha-256
    - host replication ${PATRONI_REPLICATION_USERNAME} ${DB3_WG_IP}/32 scram-sha-256
    - host all all ${DB1_WG_IP}/32 scram-sha-256
    - host all all ${DB2_WG_IP}/32 scram-sha-256
    - host all all ${DB3_WG_IP}/32 scram-sha-256
    - host all all ${PROXY1_WG_IP}/32 scram-sha-256
    - host all all 127.0.0.1/32 scram-sha-256
  parameters:
    unix_socket_directories: /var/run/postgresql

tags:
  nofailover: false
  noloadbalance: false
  clonefrom: false
  nosync: false
EOF

  docker compose up -d --build
}

deploy_proxy_stack() {
  log "Deploying proxy stack"
  mkdir -p /opt/proxy-stack/{pgadmin-data}
  cd /opt/proxy-stack

  cat > haproxy.cfg <<'EOF'
global
    log stdout format raw local0

defaults
    log global
    mode tcp
    timeout connect 5s
    timeout client  1m
    timeout server  1m

frontend postgres_write
    bind *:5432
    default_backend patroni_primary

backend patroni_primary
    mode tcp
    option httpchk GET /primary
    http-check expect status 200
    server db1 db1:5432 check port 8008
    server db2 db2:5432 check port 8008
    server db3 db3:5432 check port 8008
EOF

  cat > Caddyfile <<EOF
{
    email ${CADDY_EMAIL}
}

${PGADMIN_DOMAIN} {
    basic_auth {
        ${PGADMIN_BASIC_AUTH_USER} ${PGADMIN_BASIC_AUTH_HASH}
    }
    reverse_proxy 127.0.0.1:8080
}

${API_DOMAIN} {
    basic_auth {
        ${API_BASIC_AUTH_USER} ${API_BASIC_AUTH_HASH}
    }
    reverse_proxy 127.0.0.1:3000
}
EOF

  cat > compose.yaml <<EOF
services:
  haproxy:
    image: haproxy:3.0
    container_name: pg-haproxy
    restart: unless-stopped
    network_mode: host
    volumes:
      - ./haproxy.cfg:/usr/local/etc/haproxy/haproxy.cfg:ro

  pgadmin:
    image: dpage/pgadmin4:9.13
    container_name: pgadmin
    restart: unless-stopped
    environment:
      PGADMIN_DEFAULT_EMAIL: ${PGADMIN_DEFAULT_EMAIL}
      PGADMIN_DEFAULT_PASSWORD: ${PGADMIN_DEFAULT_PASSWORD}
    ports:
      - "127.0.0.1:8080:80"
    volumes:
      - ./pgadmin-data:/var/lib/pgadmin

  postgrest:
    image: postgrest/postgrest:v12.2.12
    container_name: postgrest
    restart: unless-stopped
    environment:
      PGRST_DB_URI: ${POSTGREST_DB_URI}
      PGRST_DB_SCHEMAS: ${POSTGREST_DB_SCHEMAS}
      PGRST_DB_ANON_ROLE: ${POSTGREST_DB_ANON_ROLE}
      PGRST_SERVER_HOST: 0.0.0.0
      PGRST_SERVER_PORT: 3000
    ports:
      - "127.0.0.1:3000:3000"

  caddy:
    image: caddy:2
    container_name: caddy
    restart: unless-stopped
    ports:
      - "80:80"
      - "443:443"
      - "443:443/udp"
    volumes:
      - ./Caddyfile:/etc/caddy/Caddyfile:ro
      - caddy_data:/data
      - caddy_config:/config

volumes:
  caddy_data:
  caddy_config:
EOF

  docker compose up -d
}

show_next_steps_db() {
  cat <<EOF

DB node install complete.

After all three DB nodes are up, verify cluster:
  docker exec -it ${NODE_NAME}-patroni patronictl -c /etc/patroni/patroni.yml list

After proxy is up, initialize on the current primary:
  psql -h ${PROXY1_PUBLIC_IP} -p 5432 -U ${PATRONI_SUPERUSER_USERNAME}
  CREATE DATABASE ${PATRONI_DB};
  \\c ${PATRONI_DB}
  CREATE EXTENSION IF NOT EXISTS pg_cron;
EOF
}

show_next_steps_proxy() {
  cat <<EOF

Proxy install complete.

Public endpoints:
  PostgreSQL via HAProxy: ${PROXY1_PUBLIC_IP}:5432
  pgAdmin: https://${PGADMIN_DOMAIN}
  API: https://${API_DOMAIN}
EOF
}

main() {
  hostnamectl set-hostname "$HOSTNAME_FQDN" || true
  install_packages
  configure_hosts
  write_wireguard

  if [[ "$ROLE" == "db" ]]; then
    configure_firewall_db
    deploy_db_stack
    show_next_steps_db
  else
    configure_firewall_proxy
    deploy_proxy_stack
    show_next_steps_proxy
  fi
}

main "$@"