#!/usr/bin/env python3
"""Submit Z-Image-Turbo workflow to ComfyUI on BC-250 via API."""
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}")
    print(f"$ {cmd[:200]}...")
    _, 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

# ComfyUI API prompt workflow for Z-Image-Turbo
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 in the foreground reflecting the sky, photorealistic, 8k, detailed",
                "clip": ["2", 0]
            }
        },
        "5": {
            "class_type": "EmptyZImageLatentImage //ZImagePowerNodes",
            "inputs": {
                "landscape": True,
                "ratio": "16:9  (widescreen)",
                "size": "medium (recommended)",
                "batch_size": 1
            }
        },
        "6": {
            "class_type": "ZSamplerTurbo //ZImagePowerNodes",
            "inputs": {
                "model": ["1", 0],
                "positive": ["4", 0],
                "latent_input": ["5", 0],
                "seed": 42,
                "steps": 8,
                "denoise": 1.0,
                "initial_noise_calibration": "off",
                "lowres_bias": False
            }
        },
        "7": {
            "class_type": "VAEDecode",
            "inputs": {
                "samples": ["6", 0],
                "vae": ["3", 0]
            }
        },
        "8": {
            "class_type": "SaveImage",
            "inputs": {
                "images": ["7", 0],
                "filename_prefix": "ZImageTurbo_BC250_test"
            }
        }
    }
}

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

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

# Monitor the queue and wait for completion
time.sleep(5)
run("bash -c 'curl -s http://localhost:8188/queue'",
    desc="Check queue status")

# Wait and check ComfyUI log for progress
for i in range(30):
    time.sleep(10)
    rc, out, _ = run(f"bash -c 'tail -20 /home/fabian/comfyui.log 2>/dev/null'",
        desc=f"ComfyUI log check {i+1}")
    if any(kw in out for kw in ['Prompt executed', 'SaveImage', 'output images']):
        print("\n  IMAGE GENERATION COMPLETE!")
        break
    if 'error' in out.lower() or 'Error' in out:
        print("\n  ERROR DETECTED — checking full log...")
        run("bash -c 'tail -50 /home/fabian/comfyui.log'", desc="Full error log")
        break

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

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