#!/bin/bash
# Generate a single wallpaper (any series)
# Usage: ./generate-wallpaper.sh "<prompt>" <output_file>

set -e

COMFYUI="http://192.168.1.57:8188"
PROMPT="${1:?Prompt required}"
OUTPUT="${2:?Output file required}"

echo "🎨 Generating Wallpaper"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "Prompt: ${PROMPT:0:100}..."
echo "Output: $OUTPUT"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

WORKFLOW=$(cat <<EOF
{
  "prompt": {
    "1": {
      "inputs": {
        "unet_name": "z_image_turbo_bf16.safetensors",
        "weight_dtype": "default"
      },
      "class_type": "UNETLoader"
    },
    "2": {
      "inputs": {
        "clip_name": "qwen_3_4b.safetensors",
        "type": "qwen_image"
      },
      "class_type": "CLIPLoader"
    },
    "3": {
      "inputs": {
        "vae_name": "ae.safetensors"
      },
      "class_type": "VAELoader"
    },
    "4": {
      "inputs": {
        "text": "$PROMPT",
        "clip": ["2", 0]
      },
      "class_type": "CLIPTextEncode"
    },
    "5": {
      "inputs": {
        "text": "low quality, blurry, distorted, ugly, text, watermark",
        "clip": ["2", 0]
      },
      "class_type": "CLIPTextEncode"
    },
    "6": {
      "inputs": {
        "width": 2560,
        "height": 1440,
        "batch_size": 1
      },
      "class_type": "EmptyLatentImage"
    },
    "7": {
      "inputs": {
        "seed": $(python3 -c "import random; print(random.randint(1, 2**32-1))"),
        "steps": 8,
        "cfg": 1.0,
        "sampler_name": "euler",
        "scheduler": "sgm_uniform",
        "denoise": 1,
        "model": ["1", 0],
        "positive": ["4", 0],
        "negative": ["5", 0],
        "latent_image": ["6", 0]
      },
      "class_type": "KSampler"
    },
    "8": {
      "inputs": {
        "samples": ["7", 0],
        "vae": ["3", 0]
      },
      "class_type": "VAEDecode"
    },
    "9": {
      "inputs": {
        "filename_prefix": "wp/wp_$(date +%Y%m%d)",
        "images": ["8", 0]
      },
      "class_type": "SaveImage"
    }
  }
}
EOF
)

echo "🚀 Queuing generation (2560x1440, 8 steps)..."
RESULT=$(curl -s -X POST "$COMFYUI/prompt" \
  -H "Content-Type: application/json" \
  -d "$WORKFLOW")

PROMPT_ID=$(echo "$RESULT" | grep -o '"prompt_id":"[^"]*"' | cut -d'"' -f4)

if [ -z "$PROMPT_ID" ]; then
    echo "❌ Failed to queue"
    echo "$RESULT"
    exit 1
fi

echo "✅ Queued: $PROMPT_ID"
echo "⏳ Waiting for generation (~15-20 seconds for wallpaper resolution)..."

# Wait for completion
for i in {1..60}; do
    sleep 2
    HISTORY=$(curl -s "$COMFYUI/history/$PROMPT_ID")

    if echo "$HISTORY" | grep -q '"status_str":"success"'; then
        FILENAME=$(echo "$HISTORY" | python3 -c "
import json, sys
try:
    history = json.load(sys.stdin)
    for prompt_id, data in history.items():
        if 'outputs' in data:
            for node_id, output in data['outputs'].items():
                if 'images' in output:
                    for img in output['images']:
                        print(img['filename'])
                        sys.exit(0)
except: pass
")

        if [ -n "$FILENAME" ]; then
            curl -s "$COMFYUI/view?filename=$FILENAME&type=output" -o "$OUTPUT"
            SIZE=$(du -h "$OUTPUT" | cut -f1)
            echo ""
            echo "✅ Complete: $OUTPUT ($SIZE)"
            exit 0
        fi
    elif echo "$HISTORY" | grep -q '"status_str":"error"'; then
        echo "❌ Generation failed"
        exit 1
    fi

    echo -n "."
done

echo ""
echo "⏱️ Timeout"
exit 1