#!/usr/bin/bash
#
# k8s/join-workers.sh — join worker nodes to the k3s cluster
#
# Joins all worker nodes defined in env.sh as k3s agents. Worker nodes
# run user workloads; control plane nodes are tainted NoSchedule.
#
set -uo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
# shellcheck source=./env.sh
source "$SCRIPT_DIR/env.sh"

echo "============================================"
echo "  Joining worker nodes to cluster"
echo "  Server: $BOOTSTRAP_NAME ($BOOTSTRAP_IP)"
echo "  Workers: ${#ALL_WNODE_NAMES[@]}"
echo "============================================"

# -------------------------------------------------------
# 1. Fetch join token from bootstrap node
# -------------------------------------------------------
echo ""
echo "--- [1/3] Fetching join token from $BOOTSTRAP_NAME ---"
JOIN_TOKEN=$(cn "$BOOTSTRAP_IP" 'cat /var/lib/rancher/k3s/server/token')
if [ -z "$JOIN_TOKEN" ] || [[ "$JOIN_TOKEN" == cat:* ]]; then
  echo "FATAL: could not fetch token. Got: ${JOIN_TOKEN:0:40}"
  exit 1
fi
echo "  Token OK (masked: ${JOIN_TOKEN:0:12}***)"

SERVER_URL="https://${BOOTSTRAP_IP}:${K3S_API_PORT}"

# -------------------------------------------------------
# 2. Install k3s-agent on each worker
# -------------------------------------------------------
for i in "${!ALL_WNODES[@]}"; do
  node_ip="${ALL_WNODES[$i]}"
  node_name="${ALL_WNODE_NAMES[$i]}"

  echo ""
  echo "--- [2/3] Joining worker: $node_name ($node_ip) ---"

  # Wipe any existing k3s first
  cn "$node_ip" '
    systemctl stop k3s-agent 2>/dev/null || true
    if [ -x /usr/local/bin/k3s-agent-uninstall.sh ]; then
      /usr/local/bin/k3s-agent-uninstall.sh
    fi
    rm -rf /etc/rancher/k3s /var/lib/rancher/k3s /var/lib/kubelet /var/lib/cni
    rm -f /etc/systemd/system/k3s-agent.service
    systemctl daemon-reload
    ip link delete cni0 2>/dev/null || true
    ip link delete flannel.1 2>/dev/null || true
  ' 2>/dev/null || true

  # Install as agent
  # shellcheck disable=SC2087 # heredoc intentionally expands local config
  ssh "${SSH_OPTS[@]}" "${SSH_USER}@${node_ip}" "sudo -n bash -s" <<REMOTE_AGENT
set -euo pipefail
export INSTALL_K3S_VERSION="$K3S_VERSION"
export K3S_URL="$SERVER_URL"
export K3S_TOKEN="$JOIN_TOKEN"
export K3S_NODE_NAME="$node_name"
curl -sfL https://get.k3s.io | sh -s - agent \
  --node-name=$node_name \
  --node-ip=$node_ip
REMOTE_AGENT

  echo "  $node_name agent install submitted."
done

# -------------------------------------------------------
# 3. Wait for all workers to appear Ready
# -------------------------------------------------------
TOTAL_NODES=$(( ${#ALL_CNODES[@]} + ${#ALL_WNODES[@]} ))
echo ""
echo "--- [3/3] Waiting for all $TOTAL_NODES nodes (${#ALL_CNODES[@]} cp + ${#ALL_WNODES[@]} workers) ---"
for i in $(seq 1 60); do
  READY_NODES=$(cn "$BOOTSTRAP_IP" 'k3s kubectl get nodes --no-headers 2>/dev/null | grep -c " Ready"' 2>/dev/null || echo 0)
  if [ "$READY_NODES" -ge "$TOTAL_NODES" ]; then
    echo "  All $TOTAL_NODES nodes Ready."
    break
  fi
  echo "  ...waiting ($i/60, $READY_NODES/$TOTAL_NODES ready)"
  sleep 10
done

echo ""
echo "============================================"
echo "  Node status:"
echo "============================================"
cn "$BOOTSTRAP_IP" 'k3s kubectl get nodes -o wide'

if [ "$READY_NODES" -ge "$TOTAL_NODES" ]; then
  echo ""
  echo "============================================"
  echo "  All workers joined. Cluster fully operational."
  echo "============================================"
else
  echo ""
  echo "WARN: $READY_NODES/$TOTAL_NODES ready. Check failing nodes."
  exit 1
fi