#!/bin/bash

set -e

# Set verbosity
if [ "${DEBUG}" = true ]; then
    set -x
    KUBEADM_VERBOSE="-v=5"
else
    KUBEADM_VERBOSE="-v=3"
fi

BIN_DIR="/usr/local/bin"
SBIN_DIR="/usr/local/sbin"
COMMAND=$1

# Define global compatibility matrix
declare -A versions=(
    ["containerd"]="v2.3.2"
    ["runc"]="v1.4.3"
    ["cni"]="v1.9.1"
    ["crictl"]="v1.36.0"
    ["kubernetes"]="v1.36.2"
    ["krel"]="v0.16.2"
)

# Helper function to display usage information
helper() {
    cat <<EOF
Usage:

  ENV=... kjoin <join|help>

  You must be sudo to run this script.

Commands:

  join: Join a node to the cluster
    - Installs all necessary prerequisites, container runtime, CNI plugins, and Kubernetes binaries.
    - Then joins the node to an existing Kubernetes cluster as a worker or control-plane node.
    - Requires either KUBEADM_CONFIG, or JOIN_URL + JOIN_TOKEN (JOIN_TOKEN_CACERT_HASH optional, JOIN_ASCP optional).
    - If KUBEADM_CONFIG is set, it is used as-is via --config and JOIN_URL/JOIN_TOKEN are not required.
    - If JOIN_TOKEN_CACERT_HASH is provided, it will be used for secure join.
    - If JOIN_TOKEN_CACERT_HASH is not provided, it will join using --discovery-token-unsafe-skip-ca-verification.
    - If JOIN_ASCP=true, the node joins as a control-plane node (requires JOIN_TOKEN_CERT_KEY).
    - Example: JOIN_URL=<control-plane-endpoint>:<port> JOIN_TOKEN=<token> JOIN_TOKEN_CACERT_HASH=sha256:<hash> KUBERNETES_VERSION=v1.30.5 kjoin join
    - Example: JOIN_URL=<control-plane-endpoint>:<port> JOIN_TOKEN=<token> JOIN_TOKEN_CERT_KEY=<key> JOIN_TOKEN_CACERT_HASH=sha256:<hash> JOIN_ASCP=true KUBERNETES_VERSION=v1.30.5 kjoin join
    - Example: KUBEADM_CONFIG=/path/to/kubeadm-join-config.yaml kjoin join
    - Set INSTALL_PACKAGES=true to install required OS packages first (apt/dnf/yum auto-detected).

  help: Print this help
    - Displays this help message.
    - Example: kjoin help

Environment variables:

  +-------------------------+-------------------------------------------------------------+------------+
  | Variable                | Description                                                 | Default    |
  +-------------------------+-------------------------------------------------------------+------------+
  | KUBERNETES_VERSION      | Version of kubernetes to install.                           | v1.36.2    |
  | CONTAINERD_VERSION      | Version of container runtime containerd.                    | v2.3.2     |
  | RUNC_VERSION            | Version of runc to install.                                 | v1.4.3     |
  | CNI_VERSION             | Version of CNI plugins to install.                          | v1.9.1     |
  | CRICTL_VERSION          | Version of crictl to install.                               | see matrix |
  | KUBEADM_CONFIG          | Path to the kubeadm config file to use.                     | Not set    |
  | BIND_PORT               | Port to use for the api-server (control-plane join only).   | 6443       |
  | JOIN_TOKEN              | Token to join the cluster.                                  | Not set    |
  | JOIN_TOKEN_CACERT_HASH  | Token Certificate Authority hash to join the cluster.       | Not set    |
  | JOIN_TOKEN_CERT_KEY     | Token Certificate Key to join as control-plane.             | Not set    |
  | JOIN_URL                | URL to join the cluster.                                    | Not set    |
  | JOIN_ASCP               | Switch to join either as control plane or worker.           | false      |
  | INSTALL_PACKAGES        | Install required OS packages via apt/dnf/yum before setup.  | false      |
  | REGISTRY_HOST           | Private registry host to authenticate to (GCP only).        | Not set    |
  | DEBUG                   | Set to true for more verbosity during script execution.     | false      |
  +-------------------------+-------------------------------------------------------------+------------+

EOF
}

# Log functions
info()  { echo "[INFO] $@"; }
warn()  { echo "[WARN] $@" >&2; }
fatal() { echo "[ERROR] $@" >&2; exit 1; }

# Setup architecture
setup_arch() {
    case ${ARCH:=$(uname -m)} in
        amd64|x86_64) ARCH=amd64 ;;
        arm64|aarch64) ARCH=arm64 ;;
        *) fatal "unsupported architecture ${ARCH}" ;;
    esac
    SUFFIX=$(uname -s | tr '[:upper:]' '[:lower:]')-${ARCH}
}

# Function to get compatible components version
get_version() {
    local component=$1
    echo "${versions[$component]}"
}

setup_env() {
    # Check if running as root
    [ "$(id -u)" -eq 0 ] || fatal "You need to be root to perform this install"
    # Set default values
    KUBERNETES_VERSION=${KUBERNETES_VERSION:-$(get_version "kubernetes")}
    RELEASE_VERSION=${RELEASE_VERSION:-$(get_version "krel")}
    JOIN_ASCP=${JOIN_ASCP:-false}
    CONTAINERD_VERSION=${CONTAINERD_VERSION:-$(get_version "containerd")}
    RUNC_VERSION=${RUNC_VERSION:-$(get_version "runc")}
    CNI_VERSION=${CNI_VERSION:-$(get_version "cni")}
    CRICTL_VERSION=${CRICTL_VERSION:-$(get_version "crictl")}
    BIND_PORT=${BIND_PORT:-6443}
    INSTALL_PACKAGES=${INSTALL_PACKAGES:-false}
    REGISTRY_HOST=${REGISTRY_HOST:-}
    DEBUG=${DEBUG:-false}
}

# Install OS packages required as prerequisites (optional, gated by INSTALL_PACKAGES)
install_packages() {
    info "installing prerequisite OS packages"

    if command -v apt-get &> /dev/null; then
        apt-get update
        apt-get install -y conntrack iptables socat iproute2 vim kmod procps ebtables ethtool wget curl iputils-ping apt-transport-https ca-certificates gpg bridge-utils
    elif command -v dnf &> /dev/null; then
        dnf install -y conntrack-tools iptables socat iproute vim-enhanced kmod procps-ng ebtables ethtool wget curl iputils ca-certificates gnupg2 bridge-utils
    elif command -v yum &> /dev/null; then
        yum install -y conntrack-tools iptables socat iproute vim-enhanced kmod procps-ng ebtables ethtool wget curl iputils ca-certificates gnupg2 bridge-utils
    else
        fatal "unsupported package manager: neither apt-get, dnf nor yum found"
    fi
}

# Check if prerequisites are installed
check_prerequisites() {
    info "Checking if prerequisites are installed"

    # List of required commands
    local required_commands=("conntrack" "socat" "ip" "iptables" "modprobe" "sysctl" "systemctl" "nsenter" "ebtables" "ethtool" "wget")

    for cmd in "${required_commands[@]}"; do
        if ! command -v $cmd &> /dev/null; then
            info "$cmd is not installed. Please install it before proceeding."
            exit 1
        fi
    done
}

# Configure system settings
configure_system_settings() {
    info "Configure system settings: "

    info "  - disable swap"
    swapoff -a

    info "  - swap entry process in /etc/fstab"
    sed -i '/^[^#].*\sswap\s/s/^/#/' /etc/fstab

    info "  - enable required kernel modules"
    cat <<EOF | tee /etc/modules-load.d/k8s.conf
overlay
br_netfilter
EOF
    modprobe overlay
    modprobe br_netfilter

    info "  - forwarding IPv4 and letting iptables see bridged traffic"
    cat <<EOF | tee /etc/sysctl.d/k8s.conf
net.bridge.bridge-nf-call-iptables  = 1
net.bridge.bridge-nf-call-ip6tables = 1
net.ipv4.ip_forward                 = 1
EOF

    info "  - apply sysctl settings"
    sysctl --system
}

# Install containerd
install_containerd() {
    info "installing containerd"
    wget -q --progress=bar https://github.com/containerd/containerd/releases/download/${CONTAINERD_VERSION}/containerd-${CONTAINERD_VERSION#?}-linux-${ARCH}.tar.gz
    tar Cxzvf /usr/local containerd-${CONTAINERD_VERSION#?}-linux-${ARCH}.tar.gz
    rm -f containerd-${CONTAINERD_VERSION#?}-linux-${ARCH}.tar.gz

    mkdir -p /usr/local/lib/systemd/system/
    wget -q --progress=bar https://raw.githubusercontent.com/containerd/containerd/main/containerd.service
    mv containerd.service /usr/local/lib/systemd/system/

    info "installing runc"
    wget -q --progress=bar https://github.com/opencontainers/runc/releases/download/${RUNC_VERSION}/runc.${ARCH}
    chmod 755 runc.${ARCH}
    mv runc.${ARCH} ${SBIN_DIR}/runc

    info "installing CNI plugins"
    mkdir -p /opt/cni/bin
    wget -q --progress=bar https://github.com/containernetworking/plugins/releases/download/${CNI_VERSION}/cni-plugins-linux-${ARCH}-${CNI_VERSION}.tgz
    tar Cxzvf /opt/cni/bin cni-plugins-linux-${ARCH}-${CNI_VERSION}.tgz
    rm -f cni-plugins-linux-${ARCH}-${CNI_VERSION}.tgz

    info "configuring systemd cgroup driver in containers"
    mkdir -p /etc/containerd
    containerd config default | sed -e "s#SystemdCgroup = false#SystemdCgroup = true#g" > /etc/containerd/config.toml

    info "configuring containerd config imports"
    mkdir -p /etc/containerd/conf.d /etc/containerd/certs.d

    cat <<EOF | tee /etc/containerd/conf.d/registry.toml
[plugins."io.containerd.cri.v1.images".registry]
   config_path = "/etc/containerd/certs.d"
EOF

    sed -i '/\[Service\]/a EnvironmentFile='/etc/environment'' /usr/local/lib/systemd/system/containerd.service
    systemctl daemon-reload && systemctl enable --now containerd && systemctl restart containerd
}

# Configure containerd authentication to a private registry using the host's own service
# account token.
#
# Why the token lands in hosts.toml and not in /var/lib/kubelet/config.json: kubelet builds its
# docker keyring once at startup, so a file written later is never picked up and every token
# rotation would need a kubelet restart. containerd re-reads hosts.toml per pull, so a timer is
# enough.
#
# Why not the CRI registry config in config.toml: containerd 2.x removed registry.configs,
# registry.auths and registry.mirrors from the CRI plugin. config_path plus hosts.toml is what is
# left, and install_containerd already points containerd at /etc/containerd/certs.d.
#
# No-op unless REGISTRY_HOST is set and the GCE metadata server answers, so the function stays
# inert on any non-GCP host.
configure_registry_auth() {
    [ -n "${REGISTRY_HOST}" ] || return 0

    if ! curl -sf --max-time 5 -H 'Metadata-Flavor: Google' \
        http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/email \
        >/dev/null 2>&1; then
        warn "REGISTRY_HOST is set but the GCE metadata server is unreachable, skipping registry auth"
        return 0
    fi

    info "configuring registry authentication for ${REGISTRY_HOST}"

    mkdir -p "/etc/containerd/certs.d/${REGISTRY_HOST}"

    # Quoted heredoc: the refresher resolves REGISTRY_HOST at runtime from the unit environment,
    # so changing the host means editing one unit file, not regenerating the script.
    cat <<'EOS' > ${BIN_DIR}/kjoin-registry-auth
#!/bin/sh
# Written by kjoin. Refreshes the private registry authorization header for containerd.
set -eu

: "${REGISTRY_HOST:?REGISTRY_HOST is not set}"
DIR="/etc/containerd/certs.d/${REGISTRY_HOST}"

TOKEN=$(curl -sf --max-time 10 -H 'Metadata-Flavor: Google' \
    http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token \
    | sed -n 's/.*"access_token":"\([^"]*\)".*/\1/p')
[ -n "${TOKEN}" ] || { echo "no access token from the metadata server" >&2; exit 1; }

mkdir -p "${DIR}"
umask 077

# The temporary file lives in the same directory so the rename is atomic: containerd may read
# hosts.toml at any moment, and a half-written file would break pulls until the next tick.
cat > "${DIR}/.hosts.toml.tmp" <<TOML
server = "https://${REGISTRY_HOST}"

[host."https://${REGISTRY_HOST}"]
  capabilities = ["pull", "resolve"]

  [host."https://${REGISTRY_HOST}".header]
    authorization = "Bearer ${TOKEN}"
TOML

mv -f "${DIR}/.hosts.toml.tmp" "${DIR}/hosts.toml"
EOS
    chmod 0755 ${BIN_DIR}/kjoin-registry-auth

    cat <<EOF > /etc/systemd/system/kjoin-registry-auth.service
[Unit]
Description=Refresh private registry credentials for containerd
After=network-online.target
Wants=network-online.target

[Service]
Type=oneshot
Environment=REGISTRY_HOST=${REGISTRY_HOST}
ExecStart=${BIN_DIR}/kjoin-registry-auth
EOF

    # The token lives for an hour; refreshing twice as often leaves room for one failed tick.
    cat <<EOF > /etc/systemd/system/kjoin-registry-auth.timer
[Unit]
Description=Refresh private registry credentials for containerd every 30 minutes

[Timer]
OnBootSec=1min
OnUnitActiveSec=30min
AccuracySec=1min

[Install]
WantedBy=timers.target
EOF

    systemctl daemon-reload
    systemctl enable --now kjoin-registry-auth.timer

    # Synchronous first run: kubeadm join pulls images, and without credentials in place by then
    # the very first pull goes out anonymous and fails.
    systemctl start kjoin-registry-auth.service
}

# Install crictl
install_crictl() {
    info "installing crictl"
    wget -qO- --progress=bar "https://github.com/kubernetes-sigs/cri-tools/releases/download/${CRICTL_VERSION}/crictl-${CRICTL_VERSION}-linux-${ARCH}.tar.gz" | tar -C "${BIN_DIR}" -xz
    rm -f crictl-${CRICTL_VERSION}-linux-${ARCH}.tar.gz  # Cleanup downloaded tar.gz file
    cat <<EOF | tee /etc/crictl.yaml
runtime-endpoint: unix:///var/run/containerd/containerd.sock
image-endpoint: unix:///var/run/containerd/containerd.sock
timeout: 10
EOF
}

# Install Kubernetes binaries
install_kube_binaries() {
   info "create folder for static pod manifests"
   mkdir -p /etc/kubernetes/manifests
   install_kubeadm
   install_kubelet
   install_kubectl
}

install_kubeadm() {
  info "installing kubeadm"
  wget -q --progress=bar https://dl.k8s.io/release/${KUBERNETES_VERSION}/bin/linux/${ARCH}/kubeadm
  chmod +x kubeadm
  mv kubeadm ${BIN_DIR}/kubeadm
}

install_kubelet() {
  info "installing kubelet"
  wget -q --progress=bar https://dl.k8s.io/release/${KUBERNETES_VERSION}/bin/linux/${ARCH}/kubelet
  chmod +x kubelet
  mv kubelet ${BIN_DIR}/kubelet

  wget -qO- --progress=bar "https://raw.githubusercontent.com/kubernetes/release/${RELEASE_VERSION}/cmd/krel/templates/latest/kubelet/kubelet.service" | sed "s:/usr/bin:${BIN_DIR}:g" | tee /usr/lib/systemd/system/kubelet.service
  chmod 600 /usr/lib/systemd/system/kubelet.service

  mkdir -p /usr/lib/systemd/system/kubelet.service.d
  wget -qO- --progress=bar "https://raw.githubusercontent.com/kubernetes/release/${RELEASE_VERSION}/cmd/krel/templates/latest/kubeadm/10-kubeadm.conf" | sed "s:/usr/bin:${BIN_DIR}:g" | tee /usr/lib/systemd/system/kubelet.service.d/10-kubeadm.conf

  systemctl daemon-reload && systemctl enable --now kubelet && systemctl restart kubelet
}

install_kubectl() {
  info "installing kubectl"
  wget -q --progress=bar https://dl.k8s.io/release/${KUBERNETES_VERSION}/bin/linux/${ARCH}/kubectl
  install -o root -g root -m 0755 kubectl ${BIN_DIR}/kubectl
  rm -f kubectl  # Cleanup downloaded binary
}

# Join node to Kubernetes cluster
join_node() {
    info "Joining the node to the cluster"

    # If KUBEADM_CONFIG is provided, use it and skip JOIN_URL/JOIN_TOKEN requirements
    if [ -n "${KUBEADM_CONFIG}" ]; then
        [ -f "${KUBEADM_CONFIG}" ] || fatal "KUBEADM_CONFIG file not found: ${KUBEADM_CONFIG}"
        local KUBEADM_ARGS="--config ${KUBEADM_CONFIG}"
    else
        # Ensure required parameters are set
        [ -z "${JOIN_URL}" ] && fatal "JOIN_URL is not set"
        [ -z "${JOIN_TOKEN}" ] && fatal "JOIN_TOKEN is not set"

        # Check if JOIN_TOKEN_CACERT_HASH is provided and not empty
        if [ -n "${JOIN_TOKEN_CACERT_HASH}" ]; then
            local KUBEADM_ARGS="${JOIN_URL} --token ${JOIN_TOKEN} --discovery-token-ca-cert-hash ${JOIN_TOKEN_CACERT_HASH}"
        else
            local KUBEADM_ARGS="${JOIN_URL} --token ${JOIN_TOKEN} --discovery-token-unsafe-skip-ca-verification"
        fi

        # If the node is joining as control plane, add control plane specific arguments
        if [ "${JOIN_ASCP}" = true ]; then
            [ -z "${JOIN_TOKEN_CERT_KEY}" ] && fatal "JOIN_TOKEN_CERT_KEY is not set for control plane join"
            KUBEADM_ARGS="${KUBEADM_ARGS} --control-plane --certificate-key ${JOIN_TOKEN_CERT_KEY} --apiserver-bind-port ${BIND_PORT}"
        fi
    fi

    # Execute kubeadm join with the constructed arguments and any additional verbosity flags
    kubeadm join ${KUBEADM_ARGS} ${KUBEADM_VERBOSE}
}

# Main commands
do_kube_setup() {
    info "Prepare the machine for kubernetes"
    [ "${INSTALL_PACKAGES}" = true ] && install_packages
    check_prerequisites
    configure_system_settings
    install_containerd
    configure_registry_auth
    install_crictl
    install_kube_binaries
}

do_kube_join() {
    do_kube_setup
    info "Joining node to the control-plane"
    join_node
}

# Main script execution
setup_arch
setup_env

case ${COMMAND} in
    join)    do_kube_join && info "join completed successfully" ;;
    help)    helper ;;
    *)       helper && fatal "use command: join|help" ;;
esac