import paramiko
import sys
import time

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

target_version = sys.argv[1] if len(sys.argv) > 1 else "2.7.10"
key_path = r"C:\Users\Snarf\.ssh\id_ed25519_snarfnet"
master_ip = "140.44.4.11"

print(f"==================================================")
print(f"Starting Rancher upgrade to version {target_version}")
print(f"==================================================")

client = paramiko.SSHClient()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

def run_cmd(cmd, check=True):
    print(f"\n>>> Running: {cmd}")
    stdin, stdout, stderr = client.exec_command(f"KUBECONFIG=/etc/rancher/k3s/k3s.yaml {cmd}")
    out = stdout.read().decode('utf-8')
    err = stderr.read().decode('utf-8')
    exit_status = stdout.channel.recv_exit_status()
    if out:
        print(out.strip())
    if err:
        print(f"STDERR: {err.strip()}")
    if check and exit_status != 0:
        raise RuntimeError(f"Command failed with exit code {exit_status}")
    return out, err, exit_status

try:
    client.connect(master_ip, username="root", key_filename=key_path, timeout=15)
    
    # 1. Take snapshot
    snap_name = f"pre-rancher-{target_version}-{int(time.time())}"
    print(f"\n--- Taking etcd snapshot: {snap_name} ---")
    run_cmd(f"k3s etcd-snapshot save --name {snap_name}")
    
    # 2. Perform Helm upgrade
    print(f"\n--- Upgrading Rancher chart to {target_version} ---")
    upgrade_cmd = (
        f"helm upgrade rancher rancher-stable/rancher "
        f"--namespace cattle-system "
        f"--version {target_version} "
        f"--reuse-values"
    )
    run_cmd(upgrade_cmd)
    
    # 3. Wait for deployment rollout
    print("\n--- Waiting for deployment/rancher rollout status ---")
    run_cmd("kubectl -n cattle-system rollout status deployment/rancher --timeout=360s")
    
    # 4. Check cattle-system pods
    print("\n--- Checking cattle-system pods ---")
    run_cmd("kubectl -n cattle-system get pods -o wide")
    
    # 5. Check ingress / HTTP response
    print("\n--- Checking Rancher endpoint ---")
    run_cmd("curl -k -I -s https://aerodrome.snarfnet.net | head -n 10", check=False)
    
    client.close()
    print(f"\nSUCCESS: Rancher upgraded to {target_version} successfully!")

except Exception as ex:
    print(f"\nFAILED: {ex}")
    sys.exit(1)