"""Save a proper Z-Image-Turbo GGUF workflow as the default ComfyUI web UI workflow."""
import paramiko, json

k = paramiko.Ed25519Key.from_private_key_file(r'C:\Users\fabia\.ssh\id_ed25519')
c = paramiko.SSHClient()
c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
c.connect('192.168.178.150', username='fabian', pkey=k, timeout=10)

# ComfyUI web UI workflow format (not API format)
workflow = {
    "last_node_id": 8,
    "last_link_id": 8,
    "nodes": [
        {
            "id": 1,
            "type": "UnetLoaderGGUF",
            "pos": [100, 100],
            "size": [300, 80],
            "flags": {},
            "order": 0,
            "mode": 0,
            "outputs": [{"name": "MODEL", "type": "MODEL", "links": [1], "slot_index": 0}],
            "properties": {"Node name for S&R": "UnetLoaderGGUF"},
            "widgets_values": ["z_image_turbo-Q5_K_S.gguf"]
        },
        {
            "id": 2,
            "type": "CLIPLoaderGGUF",
            "pos": [100, 250],
            "size": [300, 80],
            "flags": {},
            "order": 1,
            "mode": 0,
            "outputs": [{"name": "CLIP", "type": "CLIP", "links": [2], "slot_index": 0}],
            "properties": {"Node name for S&R": "CLIPLoaderGGUF"},
            "widgets_values": ["Qwen3-4B.i1-Q5_K_S.gguf", "qwen_image"]
        },
        {
            "id": 3,
            "type": "VAELoader",
            "pos": [100, 400],
            "size": [300, 60],
            "flags": {},
            "order": 2,
            "mode": 0,
            "outputs": [{"name": "VAE", "type": "VAE", "links": [3], "slot_index": 0}],
            "properties": {"Node name for S&R": "VAELoader"},
            "widgets_values": ["ae.safetensors"]
        },
        {
            "id": 4,
            "type": "CLIPTextEncode",
            "pos": [500, 250],
            "size": [400, 120],
            "flags": {},
            "order": 3,
            "mode": 0,
            "inputs": [{"name": "clip", "type": "CLIP", "link": 2}],
            "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [4], "slot_index": 0}],
            "properties": {"Node name for S&R": "CLIPTextEncode"},
            "widgets_values": ["A red fox in a snowy forest, photorealistic, highly detailed"]
        },
        {
            "id": 5,
            "type": "EmptyLatentImage",
            "pos": [500, 450],
            "size": [300, 110],
            "flags": {},
            "order": 4,
            "mode": 0,
            "outputs": [{"name": "LATENT", "type": "LATENT", "links": [5], "slot_index": 0}],
            "properties": {"Node name for S&R": "EmptyLatentImage"},
            "widgets_values": [512, 512, 1]
        },
        {
            "id": 6,
            "type": "KSampler",
            "pos": [950, 100],
            "size": [320, 474],
            "flags": {},
            "order": 5,
            "mode": 0,
            "inputs": [
                {"name": "model", "type": "MODEL", "link": 1},
                {"name": "positive", "type": "CONDITIONING", "link": 4},
                {"name": "negative", "type": "CONDITIONING", "link": None},
                {"name": "latent_image", "type": "LATENT", "link": 5}
            ],
            "outputs": [{"name": "LATENT", "type": "LATENT", "links": [6], "slot_index": 0}],
            "properties": {"Node name for S&R": "KSampler"},
            "widgets_values": [42, "fixed", 8, 1.0, "euler", "simple", 1.0]
        },
        {
            "id": 7,
            "type": "VAEDecode",
            "pos": [1350, 100],
            "size": [210, 50],
            "flags": {},
            "order": 6,
            "mode": 0,
            "inputs": [
                {"name": "samples", "type": "LATENT", "link": 6},
                {"name": "vae", "type": "VAE", "link": 3}
            ],
            "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [7], "slot_index": 0}],
            "properties": {"Node name for S&R": "VAEDecode"}
        },
        {
            "id": 8,
            "type": "SaveImage",
            "pos": [1350, 250],
            "size": [320, 270],
            "flags": {},
            "order": 7,
            "mode": 0,
            "inputs": [{"name": "images", "type": "IMAGE", "link": 7}],
            "properties": {"Node name for S&R": "SaveImage"},
            "widgets_values": ["ZImageTurbo"]
        }
    ],
    "links": [
        [1, 1, 0, 6, 0, "MODEL"],
        [2, 2, 0, 4, 0, "CLIP"],
        [3, 3, 0, 7, 1, "VAE"],
        [4, 4, 0, 6, 1, "CONDITIONING"],
        [5, 5, 0, 6, 3, "LATENT"],
        [6, 6, 0, 7, 0, "LATENT"],
        [7, 7, 0, 8, 0, "IMAGE"]
    ],
    "groups": [],
    "config": {},
    "extra": {},
    "version": 0.4
}

sftp = c.open_sftp()

# Save as default workflow
def ensure_dir(sftp, path):
    try:
        sftp.stat(path)
    except FileNotFoundError:
        ensure_dir(sftp, '/'.join(path.split('/')[:-1]))
        sftp.mkdir(path)

ensure_dir(sftp, '/home/fabian/ComfyUI/user/default/comfyui')

wf_json = json.dumps(workflow, indent=2)

# Save as default workflow
with sftp.open('/home/fabian/ComfyUI/user/default/comfyui/workflow.json', 'w') as f:
    f.write(wf_json)
print("Saved default workflow: ~/ComfyUI/user/default/comfyui/workflow.json")

# Also save a loadable copy in the ComfyUI root
with sftp.open('/home/fabian/ComfyUI/z_image_turbo_workflow.json', 'w') as f:
    f.write(wf_json)
print("Saved loadable copy: ~/ComfyUI/z_image_turbo_workflow.json")

sftp.close()
c.close()
print("\nDone. Refresh ComfyUI web UI — it will load the Z-Image-Turbo GGUF workflow by default.")
print("If it still shows the old workflow, click the menu and Load the z_image_turbo_workflow.json file.")