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

########################################
# Arch Linux Kubernetes Worker Node
# Fully automated worker node preparation
# + Official Kubernetes binaries pinned to 1.34.x
# Ready for manual kubeadm join
########################################

# ---------- Config ----------
K8S_VERSION="${K8S_VERSION:-v1.34.6}"
K8S_SERIES_REGEX='^v1\.34\.[0-9]+$'
K8S_ARCH="${K8S_ARCH:-amd64}"

# Binary locations
KUBEADM_BIN="/usr/local/bin/kubeadm"
KUBECTL_BIN="/usr/local/bin/kubectl"
KUBELET_BIN="/usr/local/bin/kubelet"

# Optional output file for your later manual join command
JOIN_HINT_FILE="${JOIN_HINT_FILE:-/root/kubeadm-join-example.txt}"

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

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

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

# ---------- Helpers ----------
require_cmd() {
  command -v "$1" >/dev/null 2>&1 || die "Required command not found: $1"
}

download_k8s_binary() {
  local name="$1"
  local tmpdir
  tmpdir="$(mktemp -d)"

  curl -fsSL -o "${tmpdir}/${name}" \
    "https://dl.k8s.io/release/${K8S_VERSION}/bin/linux/${K8S_ARCH}/${name}"

  curl -fsSL -o "${tmpdir}/${name}.sha256" \
    "https://dl.k8s.io/release/${K8S_VERSION}/bin/linux/${K8S_ARCH}/${name}.sha256"

  (
    cd "${tmpdir}"
    echo "$(cat "${name}.sha256")  ${name}" | sha256sum --check --status
  ) || die "Checksum verification failed for ${name} ${K8S_VERSION}"

  install -o root -g root -m 0755 "${tmpdir}/${name}" "/usr/local/bin/${name}"
  rm -rf "${tmpdir}"
}

install_kubelet_service() {
  log "Installing kubelet systemd service"

  mkdir -p /etc/systemd/system/kubelet.service.d
  touch /etc/default/kubelet

  cat >/etc/systemd/system/kubelet.service <<'EOF'
[Unit]
Description=kubelet: The Kubernetes Node Agent
Documentation=https://kubernetes.io/docs/
After=containerd.service network-online.target
Wants=network-online.target
Requires=containerd.service

[Service]
ExecStart=/usr/local/bin/kubelet
Restart=always
StartLimitInterval=0
RestartSec=10

[Install]
WantedBy=multi-user.target
EOF

  cat >/etc/systemd/system/kubelet.service.d/10-kubeadm.conf <<'EOF'
[Service]
Environment="KUBELET_KUBECONFIG_ARGS=--bootstrap-kubeconfig=/etc/kubernetes/bootstrap-kubelet.conf --kubeconfig=/etc/kubernetes/kubelet.conf"
Environment="KUBELET_CONFIG_ARGS=--config=/var/lib/kubelet/config.yaml"
EnvironmentFile=-/var/lib/kubelet/kubeadm-flags.env
EnvironmentFile=-/etc/default/kubelet
ExecStart=
ExecStart=/usr/local/bin/kubelet $KUBELET_KUBECONFIG_ARGS $KUBELET_CONFIG_ARGS $KUBELET_KUBEADM_ARGS $KUBELET_EXTRA_ARGS
EOF
}

cleanup_old_k8s_state() {
  log "Cleaning up any previous Kubernetes worker state"

  kubeadm reset -f >/dev/null 2>&1 || true

  rm -rf /etc/cni/net.d \
         /var/lib/cni \
         /etc/kubernetes \
         /var/lib/kubelet/pki \
         /var/lib/kubelet/config.yaml \
         /var/lib/kubelet/kubeadm-flags.env

  ip link delete cni0 2>/dev/null || true
  ip link delete flannel.1 2>/dev/null || true
  ip link delete kube-ipvs0 2>/dev/null || true
}

# ---------- Root check ----------
if [[ "${EUID}" -ne 0 ]]; then
  die "Run this script as root, for example: sudo ./worker_node_prepare.sh"
fi

# ---------- Cleanup on error ----------
on_error() {
  local exit_code=$?
  warn "Script failed on line $1 with exit code ${exit_code}"
  warn "Useful diagnostics:"
  echo "  journalctl -u containerd -u kubelet -b --no-pager | tail -n 200"
  echo "  systemctl status containerd kubelet --no-pager"
  exit "${exit_code}"
}
trap 'on_error $LINENO' ERR

# ---------- Version guard ----------
[[ "${K8S_VERSION}" =~ ${K8S_SERIES_REGEX} ]] || die \
  "This worker script is intended for Kubernetes 1.34.x to match your master node. Current K8S_VERSION=${K8S_VERSION}"

# ---------- Step 1: Disable swap ----------
log "Disabling swap immediately"
swapoff -a || true

log "Disabling swap persistently in /etc/fstab"
if [[ -f /etc/fstab ]]; then
  cp /etc/fstab /etc/fstab.bak.$(date +%Y%m%d%H%M%S)
  sed -ri '/\sswap\s/s/^/# DISABLED FOR KUBERNETES: /' /etc/fstab
fi

# ---------- Step 2: Update system ----------
log "Updating package databases and system packages"
pacman -Syu --noconfirm

# ---------- Step 3: Resolve iptables conflict automatically ----------
log "Resolving iptables backend for Kubernetes"
if pacman -Q iptables >/dev/null 2>&1; then
  log "Removing legacy iptables package so iptables-nft can be installed"
  pacman -Rdd --noconfirm iptables || true
fi

# ---------- Step 4: Install required Arch packages ----------
log "Installing runtime and support packages from Arch"
pacman -S --needed --noconfirm \
  ca-certificates \
  curl \
  containerd \
  cni-plugins \
  crictl \
  ethtool \
  iptables-nft \
  conntrack-tools \
  socat \
  tar \
  gzip \
  jq \
  openssl

# ---------- Step 5: Remove Arch Kubernetes packages if present ----------
log "Removing Arch-provided kubeadm/kubectl/kubelet if present"
for pkg in kubeadm kubectl kubelet; do
  if pacman -Q "${pkg}" >/dev/null 2>&1; then
    pacman -Rdd --noconfirm "${pkg}" || true
  fi
done

# ---------- Step 6: Install pinned Kubernetes binaries ----------
log "Installing Kubernetes binaries ${K8S_VERSION}"
download_k8s_binary kubeadm
download_k8s_binary kubectl
download_k8s_binary kubelet

require_cmd "${KUBEADM_BIN}"
require_cmd "${KUBECTL_BIN}"
require_cmd "${KUBELET_BIN}"

# ---------- Step 7: Kernel modules ----------
log "Configuring required kernel modules"
cat >/etc/modules-load.d/k8s.conf <<'EOF'
overlay
br_netfilter
EOF

modprobe overlay
modprobe br_netfilter

# ---------- Step 8: Sysctl ----------
log "Configuring Kubernetes sysctl settings"
cat >/etc/sysctl.d/99-kubernetes-cri.conf <<'EOF'
net.bridge.bridge-nf-call-iptables = 1
net.bridge.bridge-nf-call-ip6tables = 1
net.ipv4.ip_forward = 1
EOF

sysctl --system

# ---------- Step 9: Configure containerd ----------
log "Configuring containerd"
mkdir -p /etc/containerd

if [[ ! -f /etc/containerd/config.toml ]]; then
  containerd config default >/etc/containerd/config.toml
else
  cp /etc/containerd/config.toml /etc/containerd/config.toml.bak.$(date +%Y%m%d%H%M%S)
fi

sed -ri 's/^\s*SystemdCgroup = false/SystemdCgroup = true/' /etc/containerd/config.toml

# ---------- Step 10: Install kubelet service ----------
install_kubelet_service

# ---------- Step 11: Clean previous worker state ----------
cleanup_old_k8s_state

# ---------- Step 12: Enable services ----------
log "Enabling and starting containerd and kubelet"
systemctl daemon-reload
systemctl enable --now containerd
systemctl enable --now kubelet

# ---------- Step 13: Wait for containerd ----------
log "Waiting for containerd to become active"
for i in {1..20}; do
  if systemctl is-active --quiet containerd; then
    break
  fi
  sleep 1
done
systemctl is-active --quiet containerd || die "containerd did not start successfully"

# ---------- Step 14: Verify pinned versions ----------
log "Verifying installed Kubernetes component versions"
KUBEADM_VERSION="$("${KUBEADM_BIN}" version -o short 2>/dev/null || true)"
KUBECTL_VERSION="$("${KUBECTL_BIN}" version --client -o json 2>/dev/null | jq -r '.clientVersion.gitVersion // empty')"
KUBELET_VERSION="$("${KUBELET_BIN}" --version 2>/dev/null | awk '{print $2}')"

[[ "${KUBEADM_VERSION}" =~ ${K8S_SERIES_REGEX} ]] || die "kubeadm version mismatch: ${KUBEADM_VERSION}"
[[ "${KUBECTL_VERSION}" =~ ${K8S_SERIES_REGEX} ]] || die "kubectl version mismatch: ${KUBECTL_VERSION}"
[[ "${KUBELET_VERSION}" =~ ${K8S_SERIES_REGEX} ]] || die "kubelet version mismatch: ${KUBELET_VERSION}"

# ---------- Step 15: Write join hint ----------
log "Writing manual join hint"
cat >"${JOIN_HINT_FILE}" <<'EOF'
Run your worker join command manually, for example:

sudo kubeadm join <CONTROL_PLANE_IP>:6443 --token <TOKEN> \
  --discovery-token-ca-cert-hash sha256:<HASH>
EOF
chmod 600 "${JOIN_HINT_FILE}"

# ---------- Final output ----------
echo
echo "Worker node preparation is complete."
echo
echo "Pinned Kubernetes version:"
echo "  ${K8S_VERSION}"
echo
echo "Installed binaries:"
echo "  ${KUBEADM_BIN}"
echo "  ${KUBECTL_BIN}"
echo "  ${KUBELET_BIN}"
echo
echo "Services:"
echo "  containerd: $(systemctl is-active containerd || true)"
echo "  kubelet:    $(systemctl is-active kubelet || true)"
echo
echo "Next step:"
echo "  Run your kubeadm join command manually on this worker."
echo
echo "Example hint saved to:"
echo "  ${JOIN_HINT_FILE}"
echo
echo "Example:"
echo "  sudo kubeadm join <CONTROL_PLANE_IP>:6443 --token <TOKEN> \\"
echo "    --discovery-token-ca-cert-hash sha256:<HASH>"
echo