#!/usr/bin/env python3
"""Test Z-Image-Turbo end-to-end on BC-250 via ComfyUI API.

1. Query available nodes from ComfyUI
2. Build a workflow using Z-Image nodes + GGUF loader
3. Submit and wait for image generation
"""
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=300, desc=""):
    if desc:
        print(f"\n{'='*60}")
        print(f"  {desc}")
        print(f"{'='*60}")
    print(f"$ {cmd}")
    _, 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) > 60:
            print(f"  ... ({len(lines)} lines, showing last 60)")
            print('\n'.join(lines[-60:]))
        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

# 1. Check ComfyUI is still running
run("bash -c 'ss -tlnp | grep 8188'", desc="Verify ComfyUI is running")

# 2. Get available node types — look for Z-Image and GGUF related nodes
run("bash -c 'curl -s http://localhost:8188/object_info 2>/dev/null | python3 -c \""
    "import sys, json; "
    "data = json.load(sys.stdin); "
    "nodes = sorted(data.keys()); "
    "z_nodes = [n for n in nodes if any(k in n.lower() for k in [\\\"zimage\\\", \\\"z_image\\\", \\\"zi_\\\", \\\"gguf\\\", \\\"unet\\\", \\\"sampler\\\", \\\"vae\\\", \\\"clip\\\", \\\"save\\\", \\\"empty\\\", \\\"latent\\\"])]; "
    "print(\\\"Relevant nodes:\\\"); "
    "[print(f\\\"  {n}\\\") for n in z_nodes]; "
    "print(f\\\"\\\\nTotal nodes: {len(nodes)}\\\"); "
    "\"'",
    desc="Query ComfyUI for available nodes")

# 3. Get detailed info on Z-Image and GGUF nodes
run("bash -c 'curl -s http://localhost:8188/object_info 2>/dev/null | python3 -c \""
    "import sys, json; "
    "data = json.load(sys.stdin); "
    "targets = [k for k in data if any(t in k.lower() for t in [\\\"zi_\\\", \\\"zimage\\\", \\\"z_image\\\", \\\"gguf\\\"])]; "
    "for name in sorted(targets): "
    "    info = data[name]; "
    "    print(f\\\"\\\\n=== {name} ===\\\"); "
    "    inp = info.get(\\\"input\\\", {}).get(\\\"required\\\", {}); "
    "    print(f\\\"  Required inputs:\\\"); "
    "    for k, v in inp.items(): "
    "        print(f\\\"    {k}: {v}\\\"); "
    "    opt = info.get(\\\"input\\\", {}).get(\\\"optional\\\", {}); "
    "    if opt: "
    "        print(f\\\"  Optional inputs:\\\"); "
    "        for k, v in opt.items(): "
    "            print(f\\\"    {k}: {v}\\\"); "
    "    out = info.get(\\\"output\\\", []); "
    "    print(f\\\"  Outputs: {out}\\\"); "
    "\"'",
    desc="Get Z-Image and GGUF node details")

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