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

script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# shellcheck source=scripts/cluster/common.sh
source "$script_dir/common.sh"

load_cluster_env
require_commands ssh scp

hosts=("$CLUSTER_NINGBO_HOST" "$CLUSTER_HONGKONG_HOST" "$CLUSTER_LOS_ANGELES_HOST")
public_endpoints=("${CLUSTER_NINGBO_HOST#*@}" "${CLUSTER_HONGKONG_HOST#*@}" "${CLUSTER_LOS_ANGELES_HOST#*@}")
wireguard_ips=(10.77.0.1 10.77.0.2 10.77.0.3)
public_keys=()

for host in "${hosts[@]}"; do
  echo "[wireguard] preparing $host"
  public_key=$(cluster_ssh "$host" 'bash -s' <<'REMOTE'
set -euo pipefail
export DEBIAN_FRONTEND=noninteractive
if ! command -v wg >/dev/null 2>&1; then
  apt-get update -qq
  apt-get install -y -qq wireguard-tools
fi
install -d -m 0700 /etc/wireguard
if [[ ! -s /etc/wireguard/privatekey ]]; then
  umask 077
  wg genkey >/etc/wireguard/privatekey
fi
chmod 0600 /etc/wireguard/privatekey
wg pubkey </etc/wireguard/privatekey >/etc/wireguard/publickey
cat /etc/wireguard/publickey
REMOTE
)
  [[ $public_key =~ ^[A-Za-z0-9+/]{43}=$ ]] || {
    echo "invalid WireGuard public key returned by $host" >&2
    exit 1
  }
  public_keys+=("$public_key")
done

for local_index in 0 1 2; do
  peer_indexes=()
  for candidate in 0 1 2; do
    [[ $candidate -eq $local_index ]] || peer_indexes+=("$candidate")
  done
  first=${peer_indexes[0]}
  second=${peer_indexes[1]}
  echo "[wireguard] configuring ${hosts[$local_index]} as ${wireguard_ips[$local_index]}"
  cluster_ssh "${hosts[$local_index]}" bash -s -- \
    "${wireguard_ips[$local_index]}" \
    "${public_keys[$first]}" "${public_endpoints[$first]}" "${wireguard_ips[$first]}" \
    "${public_keys[$second]}" "${public_endpoints[$second]}" "${wireguard_ips[$second]}" <<'REMOTE'
set -euo pipefail
local_ip=$1
peer_one_key=$2
peer_one_endpoint=$3
peer_one_ip=$4
peer_two_key=$5
peer_two_endpoint=$6
peer_two_ip=$7
private_key=$(cat /etc/wireguard/privatekey)
umask 077
temporary_config=$(mktemp /etc/wireguard/wg0.conf.XXXXXX)
cat >"$temporary_config" <<EOF
[Interface]
Address = ${local_ip}/24
ListenPort = 51820
PrivateKey = ${private_key}

[Peer]
PublicKey = ${peer_one_key}
Endpoint = ${peer_one_endpoint}:51820
AllowedIPs = ${peer_one_ip}/32
PersistentKeepalive = 25

[Peer]
PublicKey = ${peer_two_key}
Endpoint = ${peer_two_endpoint}:51820
AllowedIPs = ${peer_two_ip}/32
PersistentKeepalive = 25
EOF
chmod 0600 "$temporary_config"
mv "$temporary_config" /etc/wireguard/wg0.conf
systemctl enable wg-quick@wg0 >/dev/null
systemctl restart wg-quick@wg0
REMOTE
done

for source_index in 0 1 2; do
  for target_index in 0 1 2; do
    [[ $source_index -eq $target_index ]] && continue
    cluster_ssh "${hosts[$source_index]}" \
      "ping -c 3 -W 2 ${wireguard_ips[$target_index]} >/dev/null"
  done
done

for host in "${hosts[@]}"; do
  latest_handshake=$(cluster_ssh "$host" \
    "wg show wg0 latest-handshakes | awk '\$2 > 0 { count++ } END { print count + 0 }'")
  [[ $latest_handshake == 2 ]] || {
    echo "WireGuard handshake verification failed on $host" >&2
    exit 1
  }
done

echo 'wireguard_bootstrap=PASS peers=3'