import paramiko, sys, time

version = sys.argv[1] if len(sys.argv) > 1 else "v1.27.16+k3s1"

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect('140.44.4.11', username='root', key_filename=r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet', timeout=10)

py_code = f"""
from awx.main.models import Project, JobTemplate, Job
import time, json

# 1. Project Update
p = Project.objects.get(name='Snarfnet Infrastructure Playbooks')
update = p.update()
for _ in range(15):
    update.refresh_from_db()
    if update.status in ['successful', 'failed', 'error']:
        break
    time.sleep(2)

print(f"Project update [{{update.id}}]: {{update.status}}")

# 2. Launch Job Template
jt = JobTemplate.objects.get(name='Aerodrome K3s Control Plane Upgrade')
job = jt.create_unified_job(extra_vars=json.dumps({{'target_k3s_version': '{version}'}}))
job.signal_start()
print(f"Launched Upgrade Job [{{job.id}}] targeting version '{version}': Status={{job.status}}")

for _ in range(120):
    job.refresh_from_db()
    if job.status in ['successful', 'failed', 'error', 'canceled']:
        break
    time.sleep(3)

print(f"Job [{{job.id}}] finished with Status={{job.status}}, Elapsed={{job.elapsed}}s, Failed={{job.failed}}")

print('\\n=== PLAY RECAP ===')
for e in job.job_events.filter(event='playbook_on_stats'):
    print(e.stdout)
"""

cmd = f"kubectl exec -i -n awx -c awx-task deployment/awx-task -- awx-manage shell_plus --quiet << 'EOF'\n{py_code}\nEOF"
stdin, stdout, stderr = ssh.exec_command(cmd)

for line in iter(stdout.readline, ""):
    print(line, end="")

# Check kubectl get nodes
print("\n=== CLUSTER STATUS AFTER STEP ===")
stdin, stdout, stderr = ssh.exec_command("kubectl get nodes -o wide")
print(stdout.read().decode())

ssh.close()