#!/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-worker-node.sh [preflight|network|install|verify|all]' >&2
  exit 64
}

load_cluster_env
require_commands ssh scp shasum
: "${CLUSTER_WORKER_NODE_4_HOST:?AI_GATEWAY_WORKER_NODE_4_HOST is required}"

node_name=$CLUSTER_WORKER_NODE_4_NAME
node_site=$CLUSTER_WORKER_NODE_4_SITE
node_ip=10.77.0.4
worker_wireguard_port=$CLUSTER_WORKER_NODE_4_WIREGUARD_PORT
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_WORKER_NODE_4_HOST" bash -s -- "$node_name" "$node_site" <<'REMOTE'
set -euo pipefail
node_name=$1
node_site=$2
[[ $(id -u) -eq 0 && $(uname -m) == x86_64 ]]
. /etc/os-release
[[ ($ID == debian && ${VERSION_ID%%.*} -ge 12) ||
   ($ID == ubuntu && ${VERSION_ID%%.*} -ge 24) ]]
(( $(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 'worker_node_preflight=PASS node=%s site=%s cpu=%s memory_available_mib=%s disk_available_gib=%s\n' \
  "$node_name" "$node_site" "$(nproc)" "$((available_kib / 1024))" \
  "$((available_bytes / 1024 / 1024 / 1024))"
REMOTE
}

assert_network_approved() {
  [[ $CLUSTER_WORKER_NODE_4_NETWORK_APPROVED == true ]] || {
    echo 'dedicated Worker private-network setup is blocked until the provider approves this non-proxy Kubernetes use in writing' >&2
    return 1
  }
}

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 worker_key host key

  worker_key=$(cluster_ssh "$CLUSTER_WORKER_NODE_4_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
)
  [[ $worker_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_WORKER_NODE_4_HOST" bash -s -- \
    "$worker_wireguard_port" \
    "${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
listen_port=$1
shift
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 = ${listen_port}
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 "$listen_port" 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 -- "$worker_key" "${CLUSTER_WORKER_NODE_4_HOST#*@}" \
      "$node_ip" "$worker_wireguard_port" <<'REMOTE'
set -euo pipefail
wg set wg0 peer "$1" allowed-ips "$3/32" endpoint "$2:$4" persistent-keepalive 25
wg-quick save wg0 >/dev/null
REMOTE
  done

  local handshakes_ready=false recent
  for _ in $(seq 1 30); do
    handshakes_ready=true
    for host in "${server_hosts[@]}" "$CLUSTER_WORKER_NODE_4_HOST"; do
      recent=$(cluster_ssh "$host" \
        "wg show wg0 latest-handshakes | awk '\$2 > 0 && systime()-\$2 < 180 {count++} END {print count+0}'")
      if [[ $recent != 3 ]]; then
        handshakes_ready=false
        break
      fi
    done
    [[ $handshakes_ready == true ]] && break
    sleep 2
  done
  [[ $handshakes_ready == true ]] || {
    echo 'Worker node WireGuard peers did not complete their initial handshakes within 60 seconds' >&2
    return 1
  }
  # Prime both directions after the new peer endpoint is installed. These packets are not
  # part of the measured 10-packet gate and only establish routes/NAT state.
  cluster_ssh "$CLUSTER_NINGBO_HOST" "ping -q -c 3 -W 2 $node_ip >/dev/null" || true
  cluster_ssh "$CLUSTER_WORKER_NODE_4_HOST" 'ping -q -c 3 -W 2 10.77.0.1 >/dev/null' || true
  cluster_ssh "$CLUSTER_HONGKONG_HOST" "ping -q -c 3 -W 2 $node_ip >/dev/null" || true
  cluster_ssh "$CLUSTER_WORKER_NODE_4_HOST" 'ping -q -c 3 -W 2 10.77.0.2 >/dev/null' || true
  cluster_ssh "$CLUSTER_LOS_ANGELES_HOST" "ping -q -c 3 -W 2 $node_ip >/dev/null" || true
  cluster_ssh "$CLUSTER_WORKER_NODE_4_HOST" 'ping -q -c 3 -W 2 10.77.0.3 >/dev/null' || true
  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_worker_node_4
  verify_path "$CLUSTER_WORKER_NODE_4_HOST" 10.77.0.1 80 worker_node_4_to_ningbo
  verify_path "$CLUSTER_HONGKONG_HOST" "$node_ip" 80 hongkong_to_worker_node_4
  verify_path "$CLUSTER_WORKER_NODE_4_HOST" 10.77.0.2 80 worker_node_4_to_hongkong
  verify_path "$CLUSTER_LOS_ANGELES_HOST" "$node_ip" 300 losangeles_to_worker_node_4
  verify_path "$CLUSTER_WORKER_NODE_4_HOST" 10.77.0.3 300 worker_node_4_to_losangeles
  local host recent
  for host in "$CLUSTER_NINGBO_HOST" "$CLUSTER_HONGKONG_HOST" "$CLUSTER_LOS_ANGELES_HOST" "$CLUSTER_WORKER_NODE_4_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 "worker_node_wireguard=PASS node=$node_name site=$node_site peers=4 recent_handshakes=12 routing=direct listen_port=$worker_wireguard_port"
}

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

  cluster_ssh "$CLUSTER_WORKER_NODE_4_HOST" 'install -d -m 0700 /etc/rancher/k3s'
  printf '%s\n' "$K3S_TOKEN" | cluster_ssh "$CLUSTER_WORKER_NODE_4_HOST" \
    'umask 077; cat >/etc/rancher/k3s/cluster-token; chmod 0600 /etc/rancher/k3s/cluster-token'
  cluster_scp "$k3s_binary" "$CLUSTER_WORKER_NODE_4_HOST:/usr/local/bin/k3s"
  cluster_scp "$registries_file" "$CLUSTER_WORKER_NODE_4_HOST:/etc/rancher/k3s/registries.yaml"
  cluster_ssh "$CLUSTER_WORKER_NODE_4_HOST" bash -s -- \
    "$k3s_version" "$k3s_sha256" "$node_name" "$node_site" <<'REMOTE'
set -euo pipefail
version=$1
expected_sha256=$2
node_name=$3
node_site=$4
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: ${node_name}
node-ip: 10.77.0.4
flannel-iface: wg0
node-label:
  - easyai.io/site=${node_site}
  - 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:32767 proto "$protocol" \
    comment easyai-nodeport-public-deny >/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 enable k3s-agent >/dev/null
systemctl is-active --quiet k3s-agent || systemctl start k3s-agent
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 ]]
  cluster_ssh "$CLUSTER_NINGBO_HOST" 'bash -s' <<'REMOTE'
set -euo pipefail
kubectl='k3s kubectl'
$kubectl label node easyai-ningbo easyai.io/worker=false --overwrite >/dev/null
$kubectl label node easyai-hongkong easyai.io/worker=true --overwrite >/dev/null
$kubectl label node easyai-los-angeles easyai.io/worker=false --overwrite >/dev/null
REMOTE
  verify_agent
}

verify_agent() {
  verify_wireguard
  cluster_ssh "$CLUSTER_NINGBO_HOST" bash -s -- "$node_name" "$node_site" <<'REMOTE'
set -euo pipefail
node=$1
site=$2
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}') == "$site" ]]
[[ $($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 jsonpath='{.metadata.labels.easyai\.io/worker}') == true ]]
[[ $($kubectl get node easyai-hongkong -o jsonpath='{.metadata.labels.easyai\.io/worker}') == true ]]
[[ $($kubectl get node easyai-ningbo -o jsonpath='{.metadata.labels.easyai\.io/worker}') == 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 'worker_node_k3s_agent=PASS node=%s site=%s role=worker-only etcd_member=false\n' "$node" "$site"
REMOTE
}

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