#!/usr/bin/env python3
"""Submit Z-Image-Turbo workflow using standard KSampler to ComfyUI on BC-250."""
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=600, 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) > 50:
            print(f"  ... ({len(lines)} lines, showing last 50)")
            print('\n'.join(lines[-50:]))
        else:
            print(out.strip())
    if err.strip():
        lines = err.strip().split('\n')
        show = lines[-20:] if len(lines) > 20 else lines
        print(f"STDERR: {chr(10).join(show)}")
    print(f"  Exit code: {rc}")
    return rc, out, err

# Z-Image-Turbo workflow using standard KSampler
# Turbo models: low steps (8), low/zero CFG (1.0 with cfg_pp or euler works)
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 illuminating snow-capped peaks, crystal clear lake reflecting the sky, photorealistic",
                "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_test"
            }
        }
    }
}

# Write workflow
workflow_json = json.dumps(workflow)
sftp = ssh.open_sftp()
with sftp.open('/tmp/zimage_workflow2.json', 'w') as f:
    f.write(workflow_json)
sftp.close()

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

response = {}
try:
    response = json.loads(out.strip())
except:
    pass

if 'error' in response:
    print(f"\nERROR: {response['error']}")
    if 'node_errors' in response:
        for node_id, errs in response['node_errors'].items():
            print(f"  Node {node_id} ({errs.get('class_type','')}): {errs.get('errors','')}")
    ssh.close()
    exit(1)

prompt_id = response.get('prompt_id', '')
print(f"\nPrompt ID: {prompt_id}")

# Monitor progress — model loading + 8 sampling steps
for i in range(60):  # up to 10 minutes
    time.sleep(10)
    rc, out, _ = run(f"bash -c 'tail -30 /home/fabian/comfyui.log 2>/dev/null'",
        desc=f"Progress {i+1} ({(i+1)*10}s)")
    
    if 'Prompt executed in' in out:
        print("\n  IMAGE GENERATION COMPLETE!")
        break
    if 'Exception' in out or 'Traceback' in out:
        print("\n  ERROR during generation!")
        run("bash -c 'tail -80 /home/fabian/comfyui.log'", desc="Error details")
        break

# Check output files
run("bash -c 'ls -la ~/ComfyUI/output/ 2>/dev/null'",
    desc="Output directory")

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