#!/usr/bin/env python3
"""Submit Z-Image-Turbo workflow and monitor CPU/threading."""
import paramiko
import json
import time

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect('192.168.178.150', username='fabian', key_filename=r'C:\Users\fabia\.ssh\id_ed25519')

def run(cmd, timeout=60, desc=""):
    if desc:
        print(f"\n{'='*60}")
        print(f"  {desc}")
        print(f"{'='*60}")
    _, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode()
    err = stderr.read().decode()
    rc = stdout.channel.recv_exit_status()
    if out.strip():
        lines = out.strip().split('\n')
        if len(lines) > 40:
            print(f"  ... ({len(lines)} lines, showing last 40)")
            print('\n'.join(lines[-40:]))
        else:
            print(out.strip())
    if err.strip():
        lines = err.strip().split('\n')[-10:]
        print(f"STDERR: {chr(10).join(lines)}")
    print(f"  Exit: {rc}")
    return rc, out, err

# Workflow JSON
workflow = {
    "prompt": {
        "1": {
            "class_type": "UnetLoaderGGUF",
            "inputs": {"unet_name": "z_image_turbo-Q5_K_S.gguf"}
        },
        "2": {
            "class_type": "CLIPLoaderGGUF",
            "inputs": {"clip_name": "Qwen3-4B.i1-Q5_K_S.gguf", "type": "qwen_image"}
        },
        "3": {
            "class_type": "VAELoader",
            "inputs": {"vae_name": "ae.safetensors"}
        },
        "4": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "text": "A majestic mountain landscape at sunset, golden light on snow peaks, crystal lake reflection, photorealistic, 8k",
                "clip": ["2", 0]
            }
        },
        "5": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": "", "clip": ["2", 0]}
        },
        "6": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": 1024, "height": 576, "batch_size": 1}
        },
        "7": {
            "class_type": "KSampler",
            "inputs": {
                "model": ["1", 0],
                "seed": 42,
                "steps": 8,
                "cfg": 1.0,
                "sampler_name": "euler",
                "scheduler": "simple",
                "positive": ["4", 0],
                "negative": ["5", 0],
                "latent_image": ["6", 0],
                "denoise": 1.0
            }
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["7", 0], "vae": ["3", 0]}
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"images": ["8", 0], "filename_prefix": "ZImageTurbo_BC250"}
        }
    }
}

# Upload and submit
sftp = ssh.open_sftp()
with sftp.open('/tmp/zimage_workflow.json', 'w') as f:
    f.write(json.dumps(workflow))
sftp.close()

rc, out, _ = run("bash -c 'curl -s -X POST http://localhost:8188/prompt "
    "-H \"Content-Type: application/json\" "
    "-d @/tmp/zimage_workflow.json'",
    desc="Submit workflow")

try:
    resp = json.loads(out.strip())
    if 'error' in resp:
        print(f"\nERROR: {resp['error']}")
        if 'node_errors' in resp:
            for nid, e in resp['node_errors'].items():
                print(f"  Node {nid}: {e.get('errors', [])}")
        ssh.close()
        exit(1)
    print(f"\nPrompt ID: {resp.get('prompt_id', 'unknown')}")
except:
    print(f"Response: {out.strip()[:500]}")

# Monitor: check CPU usage + log every 15s
print("\n  Monitoring CPU and progress...")
for i in range(80):  # up to 20 minutes
    time.sleep(15)
    
    # Get CPU usage per-thread of ComfyUI process + total system load
    rc, cpu_out, _ = run("bash -c '"
        "PID=$(pgrep -f \"python3 main.py\" | head -1); "
        "if [ -n \"$PID\" ]; then "
        "  echo \"=== Process CPU ===\"; "
        "  ps -p $PID -o pid,%cpu,%mem,nlwp --no-headers; "
        "  echo \"=== System Load ===\"; "
        "  uptime; "
        "  echo \"=== Per-Core ===\"; "
        "  mpstat -P ALL 1 1 2>/dev/null | tail -15 || cat /proc/loadavg; "
        "else echo PROCESS_DEAD; fi'")
    
    rc, log_out, _ = run("bash -c 'tail -5 /home/fabian/comfyui.log 2>/dev/null'")
    
    if 'PROCESS_DEAD' in (cpu_out or ''):
        print("\n  ComfyUI process died!")
        run("bash -c 'tail -50 /home/fabian/comfyui.log'", desc="Death log")
        break
    if 'Prompt executed in' in (log_out or ''):
        print(f"\n  IMAGE GENERATED! (at check {i+1}, ~{(i+1)*15}s)")
        run("bash -c 'tail -30 /home/fabian/comfyui.log'", desc="Completion log")
        break
    if 'Traceback' in (log_out or '') or 'Exception' in (log_out or ''):
        print("\n  ERROR!")
        run("bash -c 'tail -60 /home/fabian/comfyui.log'", desc="Error log")
        break
    
    print(f"  [{i+1}] {(i+1)*15}s elapsed...")

# Check output
run("bash -c 'ls -lah ~/ComfyUI/output/ 2>/dev/null'", desc="Output files")
run("bash -c 'tail -20 /home/fabian/comfyui.log 2>/dev/null'", desc="Final log")

ssh.close()
print("\nDone.")