import paramiko
import sys
import time

sys.stdout.reconfigure(encoding='utf-8')

MASTER_IP = '140.44.4.11'
KEY_FILE = r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet'
TARGET_VERSION = 'v1.31.4+k3s1'
ENCODED_VER = 'v1.31.4%2Bk3s1'

workers = [
    ('sninfr2lnxwrk03', '140.44.4.23'),
    ('sninfr2lnxwrk04', '140.44.4.24'),
    ('sninfr2lnxwrk05', '140.44.4.25'),
    ('sninfr2lnxwrk06', '140.44.4.26'),
    ('sninfr2lnxwrk07', '140.44.4.27')
]

master_ssh = paramiko.SSHClient()
master_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
master_ssh.connect(MASTER_IP, username='root', key_filename=KEY_FILE, timeout=10)

print("=== CHECKING CURRENT NODES ===")
_, out, _ = master_ssh.exec_command('kubectl get nodes')
print(out.read().decode())

# Ensure wrk03 is uncordoned first if stuck
master_ssh.exec_command('kubectl uncordon sninfr2lnxwrk03')

for node_name, worker_ip in workers:
    print(f"\n=======================================================")
    print(f"[*] Upgrading worker {node_name} ({worker_ip}) to {TARGET_VERSION}...")
    print(f"=======================================================")
    
    # 1. Cordon worker
    print(f"[*] Cordoning {node_name}...")
    master_ssh.exec_command(f'kubectl cordon {node_name}')
    
    # 2. Drain worker safely
    print(f"[*] Draining {node_name} (ignoring daemonsets)...")
    _, stdout, stderr = master_ssh.exec_command(
        f'kubectl drain {node_name} --ignore-daemonsets --delete-emptydir-data --force --timeout=120s'
    )
    print("Drain:", stdout.read().decode().strip()[:200])
    
    # 3. Connect to worker and update binary
    print(f"[*] Connecting directly to worker {worker_ip}...")
    w_ssh = paramiko.SSHClient()
    w_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    w_ssh.connect(worker_ip, username='root', key_filename=KEY_FILE, timeout=10)
    
    # Pre-seed debconf and ensure open-iscsi
    w_ssh.exec_command('echo "debconf debconf/frontend select Noninteractive" | debconf-set-selections; echo "linux-base linux-base/removing-running-kernel boolean false" | debconf-set-selections')
    
    dl_cmd = f"""
    set -e
    curl -sfL "https://github.com/k3s-io/k3s/releases/download/{ENCODED_VER}/k3s-arm64" -o /usr/local/bin/k3s.new
    chmod +x /usr/local/bin/k3s.new
    mv /usr/local/bin/k3s.new /usr/local/bin/k3s
    systemctl restart k3s-agent
    """
    print(f"[*] Downloading K3s binary & restarting k3s-agent...")
    stdin, stdout, stderr = w_ssh.exec_command(dl_cmd)
    stdout.channel.recv_exit_status()
    
    # Wait for k3s-agent to be active
    for _ in range(15):
        _, out, _ = w_ssh.exec_command('systemctl is-active k3s-agent')
        status = out.read().decode().strip()
        if status == 'active':
            print(f"[*] k3s-agent is active on {node_name}!")
            break
        time.sleep(2)
        
    w_ssh.close()
    
    # 4. Uncordon worker
    print(f"[*] Uncordoning {node_name}...")
    master_ssh.exec_command(f'kubectl uncordon {node_name}')
    
    # 5. Wait for node Ready in Kubernetes
    print(f"[*] Verifying node Ready state in Kubernetes...")
    for _ in range(25):
        _, out, _ = master_ssh.exec_command(
            f"kubectl get node {node_name} -o jsonpath='{{.status.conditions[?(@.type==\"Ready\")].status}}'"
        )
        if out.read().decode().strip() == 'True':
            print(f"[OK] {node_name} is fully Ready!")
            break
        time.sleep(4)
        
    # 6. Settle Longhorn storage replicas
    print(f"[*] Settling Longhorn storage replicas (15s)...")
    time.sleep(15)

print("\n=== FINAL CLUSTER NODE STATUS ===")
_, out, _ = master_ssh.exec_command('kubectl get nodes')
print(out.read().decode())

print("=== LONGHORN VOLUME HEALTH ===")
_, out, _ = master_ssh.exec_command('kubectl get volumes.longhorn.io -n longhorn-system')
print(out.read().decode())

master_ssh.close()
print("\n[SUCCESS] ALL AERODROME WORKERS UPGRADED TO v1.31.4+k3s1!")