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

########################################
# Arch Linux RKE2 Worker Node Prep
#
# WHAT THIS SCRIPT DOES:
# - Disables swap
# - Installs required Arch packages
# - Configures kernel modules and sysctl for Kubernetes
# - Configures NetworkManager to ignore CNI interfaces
# - Disables host nftables service (required for RKE2 networking)
# - Installs RKE2 agent (worker)
# - Prepares the node for manual join (DOES NOT join automatically)
#
# ======================================
# HOW TO USE THIS SCRIPT
# ======================================
#
# 1. Run this script on the worker node:
#
#    chmod +x worker_node_install.sh
#    sudo ./worker_node_install.sh
#
#    (Optional: set a custom node name)
#    sudo WORKER_NODE_NAME='worker-1' ./worker_node_install.sh
#
# --------------------------------------
# 2. On the MASTER node, get the token:
#
#    sudo cat /var/lib/rancher/rke2/server/node-token
#
# --------------------------------------
# 3. On the WORKER node, edit config:
#
#    sudo nano /etc/rancher/rke2/config.yaml
#
#    Replace with:
#
#    server: https://<MASTER_IP>:9345
#    token: <NODE_TOKEN>
#    node-name: <OPTIONAL_NAME>
#
# --------------------------------------
# 4. Start the worker:
#
#    sudo systemctl enable --now rke2-agent
#
# --------------------------------------
# 5. Verify
#
#    On worker:
#      sudo systemctl status rke2-agent --no-pager
#      sudo journalctl -u rke2-agent -n 200 --no-pager
#
#    On master:
#      sudo /var/lib/rancher/rke2/bin/kubectl get nodes -o wide
#
# ======================================
# NOTES
# ======================================
# - This script does NOT join automatically (by design)
# - Safe for reuse across multiple worker nodes
# - Compatible with your working master script
########################################

RKE2_VERSION="${RKE2_VERSION:-v1.34.5+rke2r1}"
WORKER_NODE_NAME="${WORKER_NODE_NAME:-}"

RKE2_CONFIG_DIR="/etc/rancher/rke2"
RKE2_CONFIG_FILE="${RKE2_CONFIG_DIR}/config.yaml"

log() {
  echo
  echo "============================================================"
  echo "[INFO] $*"
  echo "============================================================"
}

warn() {
  echo
  echo "[WARN] $*" >&2
}

die() {
  echo
  echo "[ERROR] $*" >&2
  exit 1
}

on_error() {
  local exit_code=$?
  local line_no=$1

  warn "Script failed on line ${line_no} with exit code ${exit_code}"
  warn "Useful diagnostics:"
  echo "  sudo systemctl status rke2-agent -l --no-pager"
  echo "  sudo journalctl -u rke2-agent -n 200 --no-pager"
  echo "  sudo cat ${RKE2_CONFIG_FILE}"
  exit "${exit_code}"
}
trap 'on_error $LINENO' ERR

require_root() {
  [[ "${EUID}" -eq 0 ]] || die "Run this script as root: sudo $0"
}

disable_swap() {
  log "Disabling swap"

  swapoff -a || true

  if [[ -f /etc/fstab ]]; then
    cp /etc/fstab "/etc/fstab.bak.$(date +%Y%m%d%H%M%S)"
    sed -Ei '/^[^#].+\s+swap\s+/ s/^/# disabled-by-rke2-worker-script /' /etc/fstab
  fi
}

install_packages() {
  log "Installing required Arch packages"

  pacman -Sy --noconfirm archlinux-keyring

  if pacman -Q iptables >/dev/null 2>&1; then
    pacman -Rdd --noconfirm iptables || true
  fi

  pacman -Syu --noconfirm
  pacman -S --needed --noconfirm \
    bash-completion \
    ca-certificates \
    cni-plugins \
    conntrack-tools \
    curl \
    ethtool \
    gzip \
    iproute2 \
    iptables-nft \
    jq \
    nfs-utils \
    open-iscsi \
    openssl \
    socat \
    tar \
    unzip \
    wget
}

configure_kernel() {
  log "Configuring kernel modules and sysctl"

  cat >/etc/modules-load.d/k8s.conf <<'EOF'
overlay
br_netfilter
EOF

  modprobe overlay
  modprobe br_netfilter

  cat >/etc/sysctl.d/90-kubernetes.conf <<'EOF'
net.bridge.bridge-nf-call-iptables = 1
net.bridge.bridge-nf-call-ip6tables = 1
net.ipv4.ip_forward = 1
EOF

  sysctl --system >/dev/null
}

configure_networkmanager() {
  if systemctl is-enabled NetworkManager >/dev/null 2>&1 || systemctl is-active NetworkManager >/dev/null 2>&1; then
    log "Configuring NetworkManager to ignore CNI interfaces"

    mkdir -p /etc/NetworkManager/conf.d
    cat >/etc/NetworkManager/conf.d/rke2-cni.conf <<'EOF'
[keyfile]
unmanaged-devices=interface-name:cali*;interface-name:flannel*;interface-name:cni*;interface-name:vxlan.calico;interface-name:kube-ipvs0;interface-name:nodelocaldns;interface-name:tunl*
EOF

    systemctl restart NetworkManager
  fi

  if systemctl list-unit-files | grep -q '^nm-cloud-setup.service'; then
    systemctl disable --now nm-cloud-setup.service || true
  fi
  if systemctl list-unit-files | grep -q '^nm-cloud-setup.timer'; then
    systemctl disable --now nm-cloud-setup.timer || true
  fi
}

enable_support_services() {
  log "Enabling support services"

  systemctl enable --now iscsid.service || true

  # IMPORTANT: nftables must be disabled for RKE2 networking
  systemctl stop nftables.service >/dev/null 2>&1 || true
  systemctl disable nftables.service >/dev/null 2>&1 || true
  nft flush ruleset >/dev/null 2>&1 || true
}

install_rke2_agent() {
  log "Installing RKE2 agent ${RKE2_VERSION}"

  mkdir -p "${RKE2_CONFIG_DIR}"

  curl -sfL https://get.rke2.io | INSTALL_RKE2_TYPE=agent INSTALL_RKE2_VERSION="${RKE2_VERSION}" sh -

  mkdir -p /etc/profile.d
  cat >/etc/profile.d/rke2-path.sh <<'EOF'
export PATH=$PATH:/var/lib/rancher/rke2/bin:/usr/local/bin
EOF
}

write_config_template() {
  log "Writing worker config template"

  {
    echo "# Fill these in before starting rke2-agent:"
    echo "# server: https://YOUR_MASTER_IP:9345"
    echo "# token: YOUR_NODE_TOKEN"
    if [[ -n "${WORKER_NODE_NAME}" ]]; then
      echo "node-name: ${WORKER_NODE_NAME}"
    else
      echo "# node-name: optional-custom-node-name"
    fi
  } > "${RKE2_CONFIG_FILE}"

  chmod 600 "${RKE2_CONFIG_FILE}"
}

disable_agent_until_manual_join() {
  log "Leaving rke2-agent disabled until manual join"

  systemctl daemon-reload
  systemctl disable --now rke2-agent.service >/dev/null 2>&1 || true
}

print_summary() {
  log "Worker node preparation complete"

  echo "RKE2 version: ${RKE2_VERSION}"
  echo "Config file:   ${RKE2_CONFIG_FILE}"
  echo "Node name:     ${WORKER_NODE_NAME:-<default hostname>}"
  echo
  echo "Node is ready. Follow instructions above to join cluster."
}

main() {
  require_root
  disable_swap
  install_packages
  configure_kernel
  configure_networkmanager
  enable_support_services
  install_rke2_agent
  write_config_template
  disable_agent_until_manual_join
  print_summary
}

main "$@"