import paramiko
import sys
import time

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

if len(sys.argv) < 2:
    print("Usage: python scripts/step_rancher.py <version>")
    sys.exit(1)

target_version = sys.argv[1]
key_path = r"C:\Users\Snarf\.ssh\id_ed25519_snarfnet"
master_ip = "140.44.4.11"

print("==================================================")
print(f"Rancher Upgrade Hop: Target Version {target_version}")
print("==================================================")

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.replace('.', '-')}-{int(time.time())}"
    print(f"\n--- 1. Taking etcd snapshot: {snap_name} ---")
    run_cmd(f"k3s etcd-snapshot save --name {snap_name}")
    
    # 2. Pull chart and patch kubeVersion
    print(f"\n--- 2. Pulling and preparing chart {target_version} ---")
    setup_cmd = f"""
rm -rf /tmp/rancher-hop && mkdir -p /tmp/rancher-hop && cd /tmp/rancher-hop
helm pull rancher-stable/rancher --version {target_version} --untar
sed -i '/kubeVersion:/d' /tmp/rancher-hop/rancher/Chart.yaml
"""
    run_cmd(setup_cmd)
    
    # 3. Perform Helm upgrade
    print(f"\n--- 3. Executing Helm upgrade to {target_version} ---")
    upgrade_cmd = (
        f"helm upgrade rancher /tmp/rancher-hop/rancher "
        f"--namespace cattle-system "
        f"--set hostname=aerodrome.snarfnet.net "
        f"--set ingress.tls.source=secret"
    )
    run_cmd(upgrade_cmd)
    
    # 4. Wait for deployment rollout
    print("\n--- 4. Waiting for deployment/rancher rollout status (timeout 600s) ---")
    run_cmd("kubectl -n cattle-system rollout status deployment/rancher --timeout=600s")
    
    # 5. Check cattle-system pods
    print("\n--- 5. Checking cattle-system pods ---")
    run_cmd("kubectl -n cattle-system get pods -o wide")
    
    # 6. Check webhook rollout if webhook deployment exists
    print("\n--- 6. Checking rancher-webhook status ---")
    run_cmd("kubectl -n cattle-system get pods -l app=rancher-webhook", check=False)
    
    # 7. Check ingress / HTTP response
    print("\n--- 7. Checking Rancher endpoint ---")
    time.sleep(5)
    run_cmd("curl -k -I -s https://aerodrome.snarfnet.net | head -n 10", check=False)
    
    # 8. Check helm release status
    print("\n--- 8. Helm release status ---")
    run_cmd("helm list -n cattle-system")
    
    client.close()
    print(f"\n==================================================")
    print(f"SUCCESS: Rancher upgraded to {target_version} successfully!")
    print(f"==================================================")

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