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

# Auto-updates HAProxy on proxy1 when adding/removing PostgreSQL nodes.
#
# Assumptions:
# - SSH access works from the admin machine to db/proxy hosts
# - Proxy stack lives at /opt/proxy-stack on proxy1
# - HAProxy config file is /opt/proxy-stack/haproxy.cfg
# - DB stack lives at /opt/db-stack on DB nodes
# - Hostnames db1/db2/db3/... and proxy1 are resolvable from the admin machine
#
# Usage examples:
#   ./cluster-node-haproxy.sh add-replica db4 10.20.0.14 root@db4.example.com root@proxy1.example.com
#   ./cluster-node-haproxy.sh remove-replica db4 root@db4.example.com root@proxy1.example.com
#   ./cluster-node-haproxy.sh switchover db2 root@db1.example.com
#   ./cluster-node-haproxy.sh list-haproxy root@proxy1.example.com

ADMIN_SSH_OPTS="${ADMIN_SSH_OPTS:--o BatchMode=yes}"
DB_STACK_DIR="${DB_STACK_DIR:-/opt/db-stack}"
PROXY_STACK_DIR="${PROXY_STACK_DIR:-/opt/proxy-stack}"
PATRONI_CONFIG_PATH="${PATRONI_CONFIG_PATH:-/etc/patroni/patroni.yml}"

usage() {
  cat <<EOF
Usage:
  $0 add-replica <new_name> <new_wg_ip> <new_host_ssh> <proxy_host_ssh>
  $0 remove-replica <old_name> <old_host_ssh> <proxy_host_ssh>
  $0 switchover <candidate_name> <admin_db_host_ssh>
  $0 list-haproxy <proxy_host_ssh>

Env:
  ADMIN_SSH_OPTS
  DB_STACK_DIR
  PROXY_STACK_DIR
  PATRONI_CONFIG_PATH
EOF
}

run_remote() {
  local host="$1"
  shift
  ssh ${ADMIN_SSH_OPTS} "${host}" "$@"
}

wait_for_patroni_member() {
  local admin_db_host_ssh="$1"
  local member="$2"
  echo "Waiting for Patroni member ${member} to appear..."
  for _ in $(seq 1 60); do
    if run_remote "${admin_db_host_ssh}" \
      "docker ps --format '{{.Names}}' | grep -q patroni && docker exec \$(docker ps --format '{{.Names}}' | grep patroni | head -n1) patronictl -c ${PATRONI_CONFIG_PATH} list" \
      | grep -qE "^[[:space:]]*${member}[[:space:]]"; then
      echo "Member ${member} found."
      return 0
    fi
    sleep 5
  done
  echo "Timed out waiting for Patroni member ${member}" >&2
  return 1
}

haproxy_add_server() {
  local proxy_host_ssh="$1"
  local node_name="$2"
  local node_ip="$3"

  run_remote "${proxy_host_ssh}" "python3 - <<'PY'
from pathlib import Path
cfg = Path('${PROXY_STACK_DIR}/haproxy.cfg')
text = cfg.read_text()
line = f'    server ${node_name} ${node_name}:5432 check port 8008'
if line in text:
    print('HAProxy already contains ${node_name}')
    raise SystemExit(0)
marker = 'backend patroni_primary\\n'
if marker not in text:
    raise SystemExit('backend patroni_primary block not found')
parts = text.split(marker, 1)
head, tail = parts[0], parts[1]
insert_after = ''
# append after backend header block
new_tail = tail + '\\n' if not tail.endswith('\\n') else tail
cfg.write_text(head + marker + tail.rstrip('\\n') + '\\n' + line + '\\n')
print('Added ${node_name} to HAProxy config')
PY"
}

haproxy_remove_server() {
  local proxy_host_ssh="$1"
  local node_name="$2"

  run_remote "${proxy_host_ssh}" "python3 - <<'PY'
from pathlib import Path
cfg = Path('${PROXY_STACK_DIR}/haproxy.cfg')
text = cfg.read_text().splitlines()
out = []
target = 'server ${node_name} '
removed = False
for line in text:
    if target in line:
        removed = True
        continue
    out.append(line)
cfg.write_text('\\n'.join(out) + '\\n')
print('Removed ${node_name} from HAProxy config' if removed else 'No HAProxy entry found for ${node_name}')
PY"
}

haproxy_validate_reload() {
  local proxy_host_ssh="$1"
  run_remote "${proxy_host_ssh}" "cd ${PROXY_STACK_DIR} && docker compose exec -T haproxy haproxy -c -f /usr/local/etc/haproxy/haproxy.cfg"
  run_remote "${proxy_host_ssh}" "cd ${PROXY_STACK_DIR} && docker compose restart haproxy"
  echo "HAProxy reloaded"
}

add_replica() {
  local new_name="$1"
  local new_ip="$2"
  local new_host_ssh="$3"
  local proxy_host_ssh="$4"

  echo "Starting DB stack on ${new_name}..."
  run_remote "${new_host_ssh}" "cd ${DB_STACK_DIR} && docker compose up -d --build"

  # pick db1 as admin if reachable through SSH hostname
  wait_for_patroni_member "root@db1" "${new_name}" || true

  haproxy_add_server "${proxy_host_ssh}" "${new_name}" "${new_ip}"
  haproxy_validate_reload "${proxy_host_ssh}"

  echo "Replica ${new_name} added and HAProxy updated."
}

remove_replica() {
  local old_name="$1"
  local old_host_ssh="$2"
  local proxy_host_ssh="$3"

  echo "Stopping DB stack on ${old_name}..."
  run_remote "${old_host_ssh}" "docker stop ${old_name}-patroni || true"

  haproxy_remove_server "${proxy_host_ssh}" "${old_name}"
  haproxy_validate_reload "${proxy_host_ssh}"

  echo "Replica ${old_name} removed and HAProxy updated."
}

switchover_primary() {
  local candidate="$1"
  local admin_db_host_ssh="$2"

  run_remote "${admin_db_host_ssh}" "docker exec \$(docker ps --format '{{.Names}}' | grep patroni | head -n1) patronictl -c ${PATRONI_CONFIG_PATH} switchover --candidate ${candidate} --force"
}

list_haproxy() {
  local proxy_host_ssh="$1"
  run_remote "${proxy_host_ssh}" "sed -n '/backend patroni_primary/,\$p' ${PROXY_STACK_DIR}/haproxy.cfg"
}

cmd="${1:-}"
case "${cmd}" in
  add-replica)
    [[ $# -eq 5 ]] || { usage; exit 1; }
    add_replica "$2" "$3" "$4" "$5"
    ;;
  remove-replica)
    [[ $# -eq 4 ]] || { usage; exit 1; }
    remove_replica "$2" "$3" "$4"
    ;;
  switchover)
    [[ $# -eq 3 ]] || { usage; exit 1; }
    switchover_primary "$2" "$3"
    ;;
  list-haproxy)
    [[ $# -eq 2 ]] || { usage; exit 1; }
    list_haproxy "$2"
    ;;
  *)
    usage
    exit 1
    ;;
esac