#!/bin/bash
#
# k3s Worker Node Provisioning Script
# Ensures version consistency across the cluster
#
# Usage: ./provision-k3s-worker.sh <new-node-ip> [<new-node-public-ip>]
#

set -e

# Configuration
CONTROL_PLANE_IP="10.42.1.1"
CONTROL_PLANE_PUBLIC="178.105.17.239"
NEW_NODE_IP="${1}"
NEW_NODE_PUBLIC_IP="${2:-$NEW_NODE_IP}"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

if [ -z "$NEW_NODE_IP" ]; then
  echo -e "${RED}Error: Missing required argument${NC}"
  echo "Usage: $0 <new-node-ip> [<new-node-public-ip>]"
  echo ""
  echo "Example: $0 10.42.1.6 167.233.121.122"
  exit 1
fi

echo -e "${GREEN}=== k3s Worker Node Provisioning ===${NC}"
echo "Control Plane: $CONTROL_PLANE_PUBLIC ($CONTROL_PLANE_IP)"
echo "New Node: $NEW_NODE_IP (public: $NEW_NODE_PUBLIC_IP)"
echo ""

# Step 1: Get current cluster version
echo -e "${YELLOW}[1/6]${NC} Checking current cluster version..."
CLUSTER_VERSION=$(ssh root@${CONTROL_PLANE_PUBLIC} 'k3s --version' | grep -oP 'v\K[0-9]+\.[0-9]+\.[0-9]+\+k3s[0-9]+' || echo "")

if [ -z "$CLUSTER_VERSION" ]; then
  echo -e "${RED}Error: Could not determine cluster version${NC}"
  exit 1
fi

echo -e "${GREEN}✓${NC} Cluster running k3s v${CLUSTER_VERSION}"
echo ""

# Step 2: Check Hetzner firewall
echo -e "${YELLOW}[2/6]${NC} Firewall configuration check"
echo -e "${RED}⚠  MANUAL ACTION REQUIRED${NC}"
echo "Ensure ${NEW_NODE_PUBLIC_IP} is added to Hetzner firewall 'fw-k3s'"
echo ""
echo "Command to add (if not already present):"
echo "  hcloud firewall add-rule fw-k3s --direction in --protocol tcp --source-ips ${NEW_NODE_PUBLIC_IP}/32 --port any"
echo ""
read -p "Press Enter when firewall is updated..."
echo ""

# Step 3: Verify SSH access
echo -e "${YELLOW}[3/6]${NC} Verifying SSH access to new node..."
if ! ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new root@${NEW_NODE_IP} 'echo ok' &>/dev/null; then
  echo -e "${RED}Error: Cannot reach ${NEW_NODE_IP} via SSH${NC}"
  exit 1
fi
echo -e "${GREEN}✓${NC} SSH access confirmed"
echo ""

# Step 4: Get K3S_TOKEN
echo -e "${YELLOW}[4/6]${NC} Retrieving k3s join token..."
K3S_TOKEN=$(ssh root@${CONTROL_PLANE_PUBLIC} "cat /var/lib/rancher/k3s/server/node-token")

if [ -z "$K3S_TOKEN" ]; then
  echo -e "${RED}Error: Could not retrieve K3S_TOKEN${NC}"
  exit 1
fi
echo -e "${GREEN}✓${NC} Token retrieved"
echo ""

# Step 5: Install k3s on new node with matching version
echo -e "${YELLOW}[5/6]${NC} Installing k3s v${CLUSTER_VERSION} on ${NEW_NODE_IP}..."
echo "This may take a few minutes..."

ssh root@${NEW_NODE_IP} "curl -sfL https://get.k3s.io | INSTALL_K3S_VERSION='v${CLUSTER_VERSION}' sh -s - agent \
  --server https://${CONTROL_PLANE_IP}:6443 \
  --token '${K3S_TOKEN}'"

if [ $? -ne 0 ]; then
  echo -e "${RED}Error: k3s installation failed${NC}"
  exit 1
fi

echo -e "${GREEN}✓${NC} k3s installed"
echo ""

# Step 6: Verify node joined successfully
echo -e "${YELLOW}[6/6]${NC} Verifying node joined the cluster..."
sleep 15  # Give the node time to register

echo ""
echo "Cluster nodes:"
kubectl get nodes -o wide

echo ""
echo "Node version details:"
kubectl get nodes -o custom-columns=NAME:.metadata.name,VERSION:.status.nodeInfo.kubeletVersion,CONTAINER-RUNTIME:.status.nodeInfo.containerRuntimeVersion

echo ""
echo -e "${GREEN}✓✓✓ Node provisioned successfully!${NC}"
echo ""
echo "Next steps:"
echo "  1. Label the node if needed: kubectl label nodes <node-name> node-role.kubernetes.io/worker=worker"
echo "  2. Monitor pod scheduling: kubectl get pods --all-namespaces -o wide | grep <node-name>"
echo "  3. Verify flannel annotation: kubectl get node <node-name> -o yaml | grep flannel"