import paramiko
import sys
import time

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

key_path = r"C:\Users\Snarf\.ssh\id_ed25519_snarfnet"
client = paramiko.SSHClient()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

target_node = "sninfr0lnxctl03"
target_ip = "140.44.4.43"

print(f"==================================================")
print(f"Rebooting {target_node} ({target_ip}) to load upgraded kernel...")
print(f"==================================================")

try:
    client.connect(target_ip, username="root", key_filename=key_path, timeout=5)
    print("Initiating reboot...")
    client.exec_command("reboot")
    client.close()
except Exception as e:
    print(f"Reboot trigger exception (expected): {e}")

# Wait for node to go down
print("Waiting 15 seconds for node to shut down...")
time.sleep(15)

# Poll for node to come back up
start_time = time.time()
connected = False
while time.time() - start_time < 180:
    try:
        client.connect(target_ip, username="root", key_filename=key_path, timeout=3)
        connected = True
        print(f"\n{target_node} is back UP!")
        break
    except Exception:
        sys.stdout.write(".")
        sys.stdout.flush()
        time.sleep(3)

if not connected:
    print(f"\nERROR: {target_node} failed to reconnect within 180s")
    sys.exit(1)

# Check kernel and modules
stdin, stdout, stderr = client.exec_command("uname -r; lsmod | grep veth || modprobe veth; lsmod | grep veth")
print("\nKernel & veth status:")
print(stdout.read().decode('utf-8'))
client.close()

# Check pods on Spaceport master
time.sleep(10)
client.connect("140.44.4.41", username="root", key_filename=key_path, timeout=10)
stdin, stdout, stderr = client.exec_command("k3s kubectl -n cattle-system get pods -o wide")
print("\nCattle system pods on Spaceport:")
print(stdout.read().decode('utf-8'))
client.close()

print("\nSUCCESS: sninfr0lnxctl03 rebooted and verified!")