#!/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"

action=${1:-all}
[[ $action == preflight || $action == network || $action == install ||
   $action == verify || $action == all ]] || {
  echo 'usage: add-shenzhen-worker-node.sh [preflight|network|install|verify|all]' >&2
  exit 64
}

load_cluster_env
require_commands ssh scp shasum
: "${CLUSTER_SHENZHEN_HOST:?AI_GATEWAY_SHENZHEN_HOST is required}"

node_name=easyai-shenzhen
node_ip=10.77.0.4
k3s_version=v1.36.2+k3s1
k3s_sha256=65a55ec56c24eab44383086166ec620a491952b7e23941a49ddca6e8a4c4b4de
secret_root="$cluster_root/.local-secrets/cluster"
cluster_environment="$secret_root/cluster.env"
k3s_binary="$secret_root/k3s-${k3s_version}-amd64"
registries_file="$cluster_root/deploy/kubernetes/registries.yaml"

assert_join_material() {
  assert_private_file "$cluster_environment"
  [[ -f $k3s_binary && ! -L $k3s_binary && -f $registries_file && ! -L $registries_file ]] || {
    echo 'pinned K3s binary or registry configuration is missing' >&2
    return 1
  }
  [[ $(shasum -a 256 "$k3s_binary" | awk '{print $1}') == "$k3s_sha256" ]] || {
    echo 'pinned K3s binary checksum mismatch' >&2
    return 1
  }
}

preflight_node() {
  cluster_ssh "$CLUSTER_SHENZHEN_HOST" bash -s <<'REMOTE'
set -euo pipefail
[[ $(id -u) -eq 0 && $(uname -m) == x86_64 ]]
. /etc/os-release
[[ $ID == debian && ${VERSION_ID%%.*} -ge 12 ]]
(( $(nproc) >= 4 ))
available_kib=$(awk '/MemAvailable:/ {print $2}' /proc/meminfo)
(( available_kib >= 6 * 1024 * 1024 ))
available_bytes=$(df -B1 --output=avail / | tail -n 1 | tr -d ' ')
(( available_bytes >= 20 * 1024 * 1024 * 1024 ))
timedatectl show -p NTPSynchronized --value | grep -Fxq yes
[[ ! -e /etc/rancher/k3s/config.yaml || -f /etc/rancher/k3s/config.yaml ]]
command -v ufw >/dev/null
ufw status | grep -Fqx 'Status: active'
printf 'shenzhen_worker_preflight=PASS cpu=%s memory_available_mib=%s disk_available_gib=%s\n' \
  "$(nproc)" "$((available_kib / 1024))" "$((available_bytes / 1024 / 1024 / 1024))"
REMOTE
}

configure_wireguard() {
  local -a server_hosts=("$CLUSTER_NINGBO_HOST" "$CLUSTER_HONGKONG_HOST" "$CLUSTER_LOS_ANGELES_HOST")
  local -a server_ips=(10.77.0.1 10.77.0.2 10.77.0.3)
  local -a server_endpoints=("${CLUSTER_NINGBO_HOST#*@}" "${CLUSTER_HONGKONG_HOST#*@}" "${CLUSTER_LOS_ANGELES_HOST#*@}")
  local -a server_keys=()
  local shenzhen_key host key

  shenzhen_key=$(cluster_ssh "$CLUSTER_SHENZHEN_HOST" 'bash -s' <<'REMOTE' | tail -n 1
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
)
  [[ $shenzhen_key =~ ^[A-Za-z0-9+/]{43}=$ ]]

  for host in "${server_hosts[@]}"; do
    key=$(cluster_ssh "$host" 'wg show wg0 public-key')
    [[ $key =~ ^[A-Za-z0-9+/]{43}=$ ]]
    server_keys+=("$key")
  done

  cluster_ssh "$CLUSTER_SHENZHEN_HOST" bash -s -- \
    "${server_keys[0]}" "${server_endpoints[0]}" "${server_ips[0]}" \
    "${server_keys[1]}" "${server_endpoints[1]}" "${server_ips[1]}" \
    "${server_keys[2]}" "${server_endpoints[2]}" "${server_ips[2]}" <<'REMOTE'
set -euo pipefail
private_key=$(cat /etc/wireguard/privatekey)
umask 077
temporary_config=$(mktemp /etc/wireguard/wg0.conf.XXXXXX)
cat >"$temporary_config" <<EOF
[Interface]
Address = 10.77.0.4/24
ListenPort = 51820
PrivateKey = ${private_key}

[Peer]
PublicKey = $1
Endpoint = $2:51820
AllowedIPs = $3/32
PersistentKeepalive = 25

[Peer]
PublicKey = $4
Endpoint = $5:51820
AllowedIPs = $6/32
PersistentKeepalive = 25

[Peer]
PublicKey = $7
Endpoint = $8:51820
AllowedIPs = $9/32
PersistentKeepalive = 25
EOF
chmod 0600 "$temporary_config"
mv "$temporary_config" /etc/wireguard/wg0.conf
for peer in "$2" "$5" "$8"; do
  ufw allow from "$peer" to any port 51820 proto udp comment easyai-wireguard >/dev/null
done
systemctl enable wg-quick@wg0 >/dev/null
systemctl restart wg-quick@wg0
REMOTE

  for host in "${server_hosts[@]}"; do
    cluster_ssh "$host" bash -s -- "$shenzhen_key" "${CLUSTER_SHENZHEN_HOST#*@}" "$node_ip" <<'REMOTE'
set -euo pipefail
wg set wg0 peer "$1" allowed-ips "$3/32" endpoint "$2:51820" persistent-keepalive 25
wg-quick save wg0
REMOTE
  done

  # The Hong Kong-Shenzhen public route is unstable. Keep the direct peers for
  # encrypted liveness, but route their data prefixes through Ningbo's healthy
  # WireGuard paths. AllowedIPs remain unique, so WireGuard selects the relay
  # deterministically instead of oscillating between direct and relayed paths.
  cluster_ssh "$CLUSTER_HONGKONG_HOST" bash -s -- \
    "${server_keys[0]}" "$shenzhen_key" <<'REMOTE'
set -euo pipefail
wg set wg0 peer "$1" allowed-ips 10.77.0.1/32,10.77.0.4/32
wg set wg0 peer "$2" allowed-ips 0.0.0.0/32
wg-quick save wg0
REMOTE
  cluster_ssh "$CLUSTER_SHENZHEN_HOST" bash -s -- \
    "${server_keys[0]}" "${server_keys[1]}" <<'REMOTE'
set -euo pipefail
wg set wg0 peer "$1" allowed-ips 10.77.0.1/32,10.77.0.2/32
wg set wg0 peer "$2" allowed-ips 0.0.0.0/32
wg-quick save wg0
REMOTE
  cluster_ssh "$CLUSTER_NINGBO_HOST" 'bash -s' <<'REMOTE'
set -euo pipefail
sysctl -qw net.ipv4.ip_forward=1
ensure_forward_rule() {
  if ! iptables -C FORWARD "$@" 2>/dev/null; then iptables -I FORWARD 1 "$@"; fi
}
ensure_forward_rule -i wg0 -o wg0 -s 10.77.0.2/32 -d 10.77.0.4/32 \
  -m comment --comment easyai-wireguard-hongkong-shenzhen -j ACCEPT
ensure_forward_rule -i wg0 -o wg0 -s 10.77.0.4/32 -d 10.77.0.2/32 \
  -m comment --comment easyai-wireguard-shenzhen-hongkong -j ACCEPT
netfilter-persistent save >/dev/null
REMOTE
  verify_wireguard
}

verify_path() {
  local source_host=$1 target_ip=$2 max_rtt=$3 label=$4 output loss average
  output=$(cluster_ssh "$source_host" "ping -q -i 0.2 -c 10 -W 2 $target_ip")
  loss=$(awk -F',' '/packet loss/ {gsub(/[^0-9.]/, "", $3); print $3}' <<<"$output")
  average=$(awk -F'=' '/min\/avg\/max/ {split($2, values, "/"); gsub(/ /, "", values[2]); print values[2]}' <<<"$output")
  [[ -n $loss && -n $average ]]
  awk -v loss="$loss" -v average="$average" -v max="$max_rtt" \
    'BEGIN {exit !(loss < 1 && average < max)}'
  echo "wireguard_path=$label loss_percent=$loss average_rtt_ms=$average"
}

verify_wireguard() {
  verify_path "$CLUSTER_NINGBO_HOST" "$node_ip" 80 ningbo_to_shenzhen
  verify_path "$CLUSTER_SHENZHEN_HOST" 10.77.0.1 80 shenzhen_to_ningbo
  verify_path "$CLUSTER_HONGKONG_HOST" "$node_ip" 80 hongkong_to_shenzhen
  verify_path "$CLUSTER_SHENZHEN_HOST" 10.77.0.2 80 shenzhen_to_hongkong
  verify_path "$CLUSTER_LOS_ANGELES_HOST" "$node_ip" 300 losangeles_to_shenzhen
  verify_path "$CLUSTER_SHENZHEN_HOST" 10.77.0.3 300 shenzhen_to_losangeles
  local host recent
  for host in "$CLUSTER_NINGBO_HOST" "$CLUSTER_HONGKONG_HOST" "$CLUSTER_LOS_ANGELES_HOST" "$CLUSTER_SHENZHEN_HOST"; do
    recent=$(cluster_ssh "$host" \
      "wg show wg0 latest-handshakes | awk '\$2 > 0 && systime()-\$2 < 180 {count++} END {print count+0}'")
    [[ $recent == 3 ]]
  done
  echo 'shenzhen_wireguard=PASS peers=4 recent_handshakes=12 hongkong_shenzhen_via=ningbo'
}

install_agent() {
  assert_join_material
  # shellcheck source=/dev/null
  source "$cluster_environment"
  : "${K3S_TOKEN:?}"
  [[ $K3S_TOKEN =~ ^[0-9a-f]{64}$ ]]

  cluster_ssh "$CLUSTER_SHENZHEN_HOST" 'install -d -m 0700 /etc/rancher/k3s'
  printf '%s\n' "$K3S_TOKEN" | cluster_ssh "$CLUSTER_SHENZHEN_HOST" \
    'umask 077; cat >/etc/rancher/k3s/cluster-token; chmod 0600 /etc/rancher/k3s/cluster-token'
  cluster_scp "$k3s_binary" "$CLUSTER_SHENZHEN_HOST:/usr/local/bin/k3s"
  cluster_scp "$registries_file" "$CLUSTER_SHENZHEN_HOST:/etc/rancher/k3s/registries.yaml"
  cluster_ssh "$CLUSTER_SHENZHEN_HOST" bash -s -- "$k3s_version" "$k3s_sha256" <<'REMOTE'
set -euo pipefail
version=$1
expected_sha256=$2
chmod 0755 /usr/local/bin/k3s
echo "$expected_sha256  /usr/local/bin/k3s" | sha256sum -c - >/dev/null
umask 077
cat >/etc/rancher/k3s/config.yaml <<'EOF'
server: https://10.77.0.1:6443
token-file: /etc/rancher/k3s/cluster-token
node-name: easyai-shenzhen
node-ip: 10.77.0.4
flannel-iface: wg0
node-label:
  - easyai.io/site=hongkong
  - easyai.io/workload=true
  - easyai.io/database=false
  - easyai.io/worker=true
node-taint:
  - easyai.io/worker-only=true:NoSchedule
EOF
chmod 0600 /etc/rancher/k3s/config.yaml /etc/rancher/k3s/registries.yaml
public_interface=$(ip -o route get 1.1.1.1 | awk '{for(i=1;i<=NF;i++) if($i=="dev"){print $(i+1); exit}}')
command -v ufw >/dev/null
ufw status | grep -Fqx 'Status: active'
ufw allow in on wg0 from 10.77.0.0/24 to any port 10250 proto tcp comment easyai-kubelet-private >/dev/null
ufw allow in on wg0 from 10.77.0.0/24 to any port 8472 proto udp comment easyai-flannel-private >/dev/null
ufw deny in on "$public_interface" to any port 6443 proto tcp comment easyai-k3s-api-public-deny >/dev/null
ufw deny in on "$public_interface" to any port 10250 proto tcp comment easyai-kubelet-public-deny >/dev/null
ufw deny in on "$public_interface" to any port 8472 proto udp comment easyai-flannel-public-deny >/dev/null
for protocol in tcp udp; do
  ufw deny in on "$public_interface" to any port 30000:31057 proto "$protocol" \
    comment easyai-nodeport-public-deny-low >/dev/null
  ufw deny in on "$public_interface" to any port 31059:32767 proto "$protocol" \
    comment easyai-nodeport-public-deny-high >/dev/null
done
install_script=$(mktemp)
trap 'rm -f -- "$install_script"' EXIT
curl -fsSL https://get.k3s.io -o "$install_script"
chmod 0700 "$install_script"
INSTALL_K3S_SKIP_DOWNLOAD=true INSTALL_K3S_VERSION="$version" \
  INSTALL_K3S_EXEC=agent "$install_script"
systemctl is-active --quiet k3s-agent
systemctl is-enabled --quiet k3s-agent
[[ ! -e /var/lib/rancher/k3s/server/db/etcd ]]
REMOTE

  local ready=false
  for _ in $(seq 1 90); do
    if cluster_ssh "$CLUSTER_NINGBO_HOST" \
      "k3s kubectl get node $node_name --no-headers 2>/dev/null | awk '\$2==\"Ready\" {found=1} END {exit !found}'"; then
      ready=true
      break
    fi
    sleep 2
  done
  [[ $ready == true ]]
  verify_agent
}

verify_agent() {
  verify_wireguard
  cluster_ssh "$CLUSTER_NINGBO_HOST" bash -s -- "$node_name" <<'REMOTE'
set -euo pipefail
node=$1
kubectl='k3s kubectl'
[[ $($kubectl get node "$node" -o jsonpath='{.status.conditions[?(@.type=="Ready")].status}') == True ]]
[[ $($kubectl get node "$node" -o jsonpath='{.metadata.labels.easyai\.io/site}') == hongkong ]]
[[ $($kubectl get node "$node" -o jsonpath='{.metadata.labels.easyai\.io/workload}') == true ]]
[[ $($kubectl get node "$node" -o jsonpath='{.metadata.labels.easyai\.io/database}') == false ]]
[[ $($kubectl get node "$node" -o json | jq '[.spec.taints[]? | select(.key=="easyai.io/worker-only" and .effect=="NoSchedule")] | length') -eq 1 ]]
[[ $($kubectl get pods -A --field-selector "spec.nodeName=$node" -o json |
  jq '[.items[] | select(.metadata.namespace!="kube-system")] | length') -eq 0 ]]
printf 'shenzhen_k3s_agent=PASS node=%s role=worker-only etcd_member=false\n' "$node"
REMOTE
}

case $action in
  preflight) preflight_node ;;
  network) preflight_node; configure_wireguard ;;
  install) preflight_node; install_agent ;;
  verify) preflight_node; verify_agent ;;
  all) preflight_node; configure_wireguard; install_agent ;;
esac