#!/bin/bash
set -euo pipefail

# Kubernetes Cluster Bootstrap Script
# Supports RKE2 and k3s

K8S_DISTRO="${K8S_DISTRO:-rke2}"
K8S_VERSION="${K8S_VERSION:-latest}"
NODE_TYPE="${NODE_TYPE:-server}"
MASTER_NODES="${MASTER_NODES:-}"
TOKEN="${TOKEN:-}"

log() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] $*" >&2
}

error() {
    log "ERROR: $*"
    exit 1
}

install_rke2() {
    log "Installing RKE2 ${K8S_VERSION}..."
    
    # Install RKE2
    curl -sfL https://get.rke2.io | INSTALL_RKE2_VERSION="${K8S_VERSION}" sh -
    
    # Configure RKE2
    mkdir -p /etc/rancher/rke2
    
    if [ "${NODE_TYPE}" = "server" ]; then
        cat > /etc/rancher/rke2/config.yaml <<EOF
token: ${TOKEN:-$(openssl rand -hex 32)}
cluster-cidr: "10.42.0.0/16"
service-cidr: "10.43.0.0/16"
cluster-dns: "10.43.0.10"
EOF
        
        # Enable required features
        systemctl enable rke2-server.service
        systemctl start rke2-server.service
    else
        cat > /etc/rancher/rke2/config.yaml <<EOF
server: https://${MASTER_NODES}:9345
token: ${TOKEN}
EOF
        
        systemctl enable rke2-agent.service
        systemctl start rke2-agent.service
    fi
    
    # Wait for service to be ready
    log "Waiting for RKE2 to be ready..."
    sleep 30
    
    # Install kubectl
    if [ "${NODE_TYPE}" = "server" ]; then
        mkdir -p /usr/local/bin
        cp /var/lib/rancher/rke2/bin/kubectl /usr/local/bin/kubectl
        chmod +x /usr/local/bin/kubectl
        
        # Configure kubeconfig
        mkdir -p ~/.kube
        cp /etc/rancher/rke2/rke2.yaml ~/.kube/config
        chmod 600 ~/.kube/config
    fi
}

install_k3s() {
    log "Installing k3s ${K8S_VERSION}..."
    
    if [ "${NODE_TYPE}" = "server" ]; then
        curl -sfL https://get.k3s.io | INSTALL_K3S_VERSION="${K8S_VERSION}" sh -s - server \
            --cluster-init \
            --cluster-cidr 10.42.0.0/16 \
            --service-cidr 10.43.0.0/16
        
        # Wait for k3s to be ready
        log "Waiting for k3s to be ready..."
        sleep 30
        
        # Configure kubeconfig
        mkdir -p ~/.kube
        cp /etc/rancher/k3s/k3s.yaml ~/.kube/config
        chmod 600 ~/.kube/config
    else
        curl -sfL https://get.k3s.io | INSTALL_K3S_VERSION="${K8S_VERSION}" K3S_URL=https://${MASTER_NODES}:6443 K3S_TOKEN=${TOKEN} sh -
    fi
}

setup_system() {
    log "Setting up system prerequisites..."
    
    # Disable swap
    swapoff -a
    sed -i '/ swap / s/^\(.*\)$/#\1/g' /etc/fstab
    
    # Load required kernel modules
    modprobe overlay
    modprobe br_netfilter
    
    # Configure sysctl
    cat > /etc/sysctl.d/99-kubernetes-cri.conf <<EOF
net.bridge.bridge-nf-call-iptables  = 1
net.ipv4.ip_forward                 = 1
net.bridge.bridge-nf-call-ip6tables = 1
EOF
    sysctl --system
    
    # Install required packages
    if command -v apt-get &> /dev/null; then
        apt-get update
        apt-get install -y curl wget git jq
    elif command -v yum &> /dev/null; then
        yum install -y curl wget git jq
    fi
}

install_network_plugin() {
    log "Installing network plugin (Cilium)..."
    
    kubectl apply -f https://raw.githubusercontent.com/cilium/cilium/1.14.0/install/kubernetes/quick-install.yaml
    
    log "Waiting for Cilium to be ready..."
    kubectl wait --for=condition=ready pod -l k8s-app=cilium -n kube-system --timeout=300s
}

install_storage_class() {
    log "Installing local-path storage class..."
    
    kubectl apply -f https://raw.githubusercontent.com/rancher/local-path-provisioner/v0.0.24/deploy/local-path-storage.yaml
    
    # Set as default
    kubectl patch storageclass local-path -p '{"metadata": {"annotations":{"storageclass.kubernetes.io/is-default-class":"true"}}}'
}

main() {
    log "Starting Kubernetes cluster bootstrap..."
    
    setup_system
    
    case "${K8S_DISTRO}" in
        rke2)
            install_rke2
            ;;
        k3s)
            install_k3s
            ;;
        *)
            error "Unsupported Kubernetes distribution: ${K8S_DISTRO}"
            ;;
    esac
    
    if [ "${NODE_TYPE}" = "server" ]; then
        install_network_plugin
        install_storage_class
        
        log "Kubernetes cluster bootstrap completed!"
        log "Kubeconfig location: ~/.kube/config"
        kubectl get nodes
    else
        log "Agent node setup completed!"
    fi
}

main "$@"