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())

def run_remote(ip, cmd):
    print(f"\n[{ip}] >>> Running: {cmd}")
    client.connect(ip, username="root", key_filename=key_path, timeout=15)
    stdin, stdout, stderr = client.exec_command(cmd)
    out = stdout.read().decode('utf-8')
    err = stderr.read().decode('utf-8')
    client.close()
    if out:
        print(out.strip())
    if err:
        print(f"STDERR: {err.strip()}")
    return out, err

print("==================================================")
print("1. Cleaning up Aerodrome (140.44.4.11)...")
print("==================================================")

# A. Clean up completed/failed helm-operation pods in cattle-system
cleanup_pods_cmd = (
    "KUBECONFIG=/etc/rancher/k3s/k3s.yaml kubectl -n cattle-system delete pod "
    "--field-selector=status.phase=Succeeded"
)
run_remote("140.44.4.11", cleanup_pods_cmd)

cleanup_failed_pods_cmd = (
    "KUBECONFIG=/etc/rancher/k3s/k3s.yaml kubectl -n cattle-system delete pod "
    "--field-selector=status.phase=Failed"
)
run_remote("140.44.4.11", cleanup_failed_pods_cmd)

# B. Clean up temporary files on Aerodrome
cleanup_files_cmd = "rm -rf /tmp/rancher-hop /tmp/rancher-test /tmp/rancher-2* /tmp/rancher_versions.txt"
run_remote("140.44.4.11", cleanup_files_cmd)

# C. Save final etcd snapshot on Aerodrome
snap_aero = f"final-post-upgrade-aerodrome-{int(time.time())}"
run_remote("140.44.4.11", f"k3s etcd-snapshot save --name {snap_aero}")

print("\n==================================================")
print("2. Cleaning up Spaceport (140.44.4.51)...")
print("==================================================")

# A. Clean up completed pods across namespaces on Spaceport
cleanup_spaceport_pods = (
    "KUBECONFIG=/etc/rancher/k3s/k3s.yaml kubectl delete pod -A "
    "--field-selector=status.phase=Succeeded"
)
run_remote("140.44.4.51", cleanup_spaceport_pods)

# B. Save final etcd snapshot on Spaceport
snap_space = f"final-post-upgrade-spaceport-{int(time.time())}"
run_remote("140.44.4.51", f"k3s etcd-snapshot save --name {snap_space}")

print("\n==================================================")
print("3. Verification...")
print("==================================================")

run_remote("140.44.4.11", "KUBECONFIG=/etc/rancher/k3s/k3s.yaml kubectl -n cattle-system get pods")
run_remote("140.44.4.11", "k3s etcd-snapshot list | tail -n 5")
run_remote("140.44.4.51", "k3s etcd-snapshot list | tail -n 5")

print("\n==================================================")
print("Post-upgrade cleanup completed successfully!")
print("==================================================")