#!/bin/bash
# Test image generation - v3 with CLIPLoader (not GGUF) for safetensors clip
API="http://127.0.0.1:8188"

echo "=== Queueing Z-Image Turbo generation v3 ==="
WORKFLOW='{
  "prompt": {
    "1": {
      "class_type": "UnetLoaderGGUF",
      "inputs": {
        "unet_name": "z_image_turbo-Q5_K_S.gguf"
      }
    },
    "2": {
      "class_type": "CLIPLoader",
      "inputs": {
        "clip_name": "gemma2_2b_lumina2.safetensors",
        "type": "lumina2"
      }
    },
    "3": {
      "class_type": "CLIPTextEncode",
      "inputs": {
        "text": "a beautiful mountain landscape at sunset, golden light, detailed, 4k",
        "clip": ["2", 0]
      }
    },
    "4": {
      "class_type": "CLIPTextEncode",
      "inputs": {
        "text": "blurry, ugly, distorted",
        "clip": ["2", 0]
      }
    },
    "5": {
      "class_type": "EmptyLatentImage",
      "inputs": {
        "width": 512,
        "height": 512,
        "batch_size": 1
      }
    },
    "6": {
      "class_type": "KSampler",
      "inputs": {
        "model": ["1", 0],
        "positive": ["3", 0],
        "negative": ["4", 0],
        "latent_image": ["5", 0],
        "seed": 42,
        "steps": 8,
        "cfg": 3.0,
        "sampler_name": "euler",
        "scheduler": "normal",
        "denoise": 1.0
      }
    },
    "7": {
      "class_type": "VAELoader",
      "inputs": {
        "vae_name": "ae.safetensors"
      }
    },
    "8": {
      "class_type": "VAEDecode",
      "inputs": {
        "samples": ["6", 0],
        "vae": ["7", 0]
      }
    },
    "9": {
      "class_type": "SaveImage",
      "inputs": {
        "images": ["8", 0],
        "filename_prefix": "bc250_test"
      }
    }
  }
}'

RESPONSE=$(curl -s -X POST "$API/prompt" -H "Content-Type: application/json" -d "$WORKFLOW")
echo "Response: $(echo $RESPONSE | head -c 200)"
echo ""

PROMPT_ID=$(echo "$RESPONSE" | python3.11 -c "import sys,json; print(json.load(sys.stdin).get('prompt_id','NONE'))" 2>/dev/null)
echo "Prompt ID: $PROMPT_ID"

if [ "$PROMPT_ID" = "NONE" ] || [ -z "$PROMPT_ID" ]; then
    echo "ERROR: Failed to queue!"
    echo "Full response: $RESPONSE"
    exit 1
fi

echo ""
echo "=== Waiting for generation (up to 10 min) ==="
for i in $(seq 1 120); do
    sleep 5
    STATUS=$(curl -s "$API/history/$PROMPT_ID" 2>/dev/null)
    HAS_OUTPUT=$(echo "$STATUS" | python3.11 -c "
import sys, json
data = json.load(sys.stdin)
pid = '$PROMPT_ID'
if pid in data:
    outputs = data[pid].get('outputs', {})
    status = data[pid].get('status', {})
    if status.get('status_str') == 'error':
        msgs = status.get('messages', [])
        for m in msgs:
            if isinstance(m, list) and len(m)>1 and isinstance(m[1],dict):
                em = m[1].get('exception_message','')
                if em:
                    print('ERROR:' + em[:200])
                    break
        else:
            print('ERROR:unknown')
    elif '9' in outputs:
        images = outputs['9'].get('images', [])
        if images:
            print('DONE:' + images[0].get('filename', 'unknown'))
        else:
            print('PROCESSING')
    else:
        print('PROCESSING')
else:
    print('WAITING')
" 2>/dev/null)
    echo "[$((i*5))s] $HAS_OUTPUT"
    if [[ "$HAS_OUTPUT" == DONE:* ]]; then
        FILENAME=${HAS_OUTPUT#DONE:}
        echo ""
        echo "=========================================="
        echo "  SUCCESS - Image generated!"
        echo "  File: /home/fabian/ComfyUI/output/$FILENAME"
        ls -lh "/home/fabian/ComfyUI/output/$FILENAME" 2>/dev/null
        echo "=========================================="
        exit 0
    fi
    if [[ "$HAS_OUTPUT" == ERROR:* ]]; then
        echo ""
        echo "GENERATION FAILED: $HAS_OUTPUT"
        tail -10 /home/fabian/comfyui2.log
        exit 1
    fi
done
echo "TIMEOUT"