#!/bin/bash

set -e

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

# Configuration
MASTER_NODE="turing1"
WORKER_NODES=("turing2" "turing3" "turing4")
BEELINK_NODE="beelink.lan"
MASTER_IP="192.168.222.237"
TOKEN="torino"
SSH_PASSWORD="turing"

# Function to print colored output
print_status() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Function to check if sshpass is available
check_dependencies() {
    if ! command -v sshpass &> /dev/null; then
        print_warning "sshpass not found, installing..."
        if command -v pacman &> /dev/null; then
            sudo pacman -S sshpass --noconfirm
        elif command -v apt-get &> /dev/null; then
            sudo apt-get update && sudo apt-get install -y sshpass
        else
            print_error "Cannot install sshpass automatically. Please install it manually."
            exit 1
        fi
    fi
}

# Function to upgrade master node
upgrade_master() {
    print_status "Upgrading master node: $MASTER_NODE"
    
    sshpass -p "$SSH_PASSWORD" ssh -o StrictHostKeyChecking=no root@$MASTER_NODE '
        curl -sfL https://get.k3s.io | INSTALL_K3S_SKIP_DOWNLOAD=false sh -s - \
          --write-kubeconfig-mode 644 \
          --disable servicelb \
          --token torino \
          --node-ip 192.168.222.237 \
          --disable-cloud-controller \
          --disable local-storage
    '
    
    if [ $? -eq 0 ]; then
        print_status "Master node $MASTER_NODE upgraded successfully"
    else
        print_error "Failed to upgrade master node $MASTER_NODE"
        exit 1
    fi
}

# Function to upgrade worker nodes
upgrade_worker() {
    local node=$1
    print_status "Upgrading worker node: $node"
    
    sshpass -p "$SSH_PASSWORD" ssh -o StrictHostKeyChecking=no root@$node "
        curl -sfL https://get.k3s.io | INSTALL_K3S_SKIP_DOWNLOAD=false \
          K3S_URL=https://$MASTER_IP:6443 \
          K3S_TOKEN=$TOKEN sh -
    "
    
    if [ $? -eq 0 ]; then
        print_status "Worker node $node upgraded successfully"
    else
        print_error "Failed to upgrade worker node $node"
        return 1
    fi
}

# Function to upgrade beelink node
upgrade_beelink() {
    print_status "Upgrading beelink node: $BEELINK_NODE"
    
    ssh -o PasswordAuthentication=no -o StrictHostKeyChecking=no gilgamezh@$BEELINK_NODE "
        sudo curl -sfL https://get.k3s.io | INSTALL_K3S_SKIP_DOWNLOAD=false \
          K3S_URL=https://$MASTER_IP:6443 \
          K3S_TOKEN=$TOKEN sh -
    "
    
    if [ $? -eq 0 ]; then
        print_status "Beelink node $BEELINK_NODE upgraded successfully"
    else
        print_error "Failed to upgrade beelink node $BEELINK_NODE"
        return 1
    fi
}

# Function to verify cluster health
verify_cluster() {
    print_status "Verifying cluster health..."
    
    # Wait a moment for nodes to register
    sleep 10
    
    print_status "Cluster nodes:"
    kubectl get nodes
    
    print_status "Checking for unhealthy pods..."
    unhealthy_pods=$(kubectl get pods --all-namespaces | grep -v Running | grep -v Completed | wc -l)
    
    if [ "$unhealthy_pods" -gt 1 ]; then  # Greater than 1 because header line counts
        print_warning "Found unhealthy pods:"
        kubectl get pods --all-namespaces | grep -v Running | grep -v Completed
    else
        print_status "All pods are healthy"
    fi
    
    print_status "Cluster upgrade completed successfully!"
}

# Main execution
main() {
    print_status "Starting K3s cluster upgrade..."
    print_status "This will upgrade all nodes in the TuringPi cluster"
    
    # Check dependencies
    check_dependencies
    
    # Show current cluster state
    print_status "Current cluster state:"
    kubectl get nodes
    
    read -p "Do you want to continue with the upgrade? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        print_status "Upgrade cancelled"
        exit 0
    fi
    
    # Upgrade master node first
    upgrade_master
    
    # Wait a bit for master to stabilize
    sleep 15
    
    # Upgrade worker nodes
    failed_workers=0
    for worker in "${WORKER_NODES[@]}"; do
        if ! upgrade_worker "$worker"; then
            ((failed_workers++))
        fi
        sleep 5  # Brief pause between worker upgrades
    done
    
    # Upgrade beelink node
    if ! upgrade_beelink; then
        print_warning "Beelink node upgrade failed, but continuing..."
    fi
    
    # Verify cluster health
    verify_cluster
    
    if [ $failed_workers -gt 0 ]; then
        print_warning "Upgrade completed with $failed_workers failed worker node(s)"
        exit 1
    else
        print_status "All nodes upgraded successfully!"
    fi
}

# Run main function
main "$@"