#!/bin/bash
# Wallpaper Image Generator — Cloudflare AI REST API
# Calls CF Workers AI directly (no worker needed), supports wallpaper resolution.
#
# Usage: ./cf-generate-wp.sh "<prompt>" <output_dir> <filename_prefix> [steps] [width] [height]
# Env:   CF_IMG_API_KEY   — Cloudflare API token (Workers AI: Read)
#        CF_ACCOUNT_ID    — Cloudflare account ID
#
# Defaults: 8 steps, 1920x1080

set -e

CF_API_TOKEN="${CF_API_TOKEN:?CF_API_TOKEN env var required}"
CF_ACCOUNT_ID="${CF_ACCOUNT_ID:?CF_ACCOUNT_ID env var required}"
MODEL="${CF_WP_MODEL:-@cf/black-forest-labs/flux-1-schnell}"
API_URL="https://api.cloudflare.com/client/v4/accounts/$CF_ACCOUNT_ID/ai/run/$MODEL"

PROMPT="${1:?Prompt required}"
OUTPUT_DIR="${2:?Output dir required}"
FILENAME_PREFIX="${3:?Filename prefix required}"
STEPS="${4:-8}"
WIDTH="${5:-1920}"
HEIGHT="${6:-1080}"

mkdir -p "$OUTPUT_DIR"

SEED=$(python3 -c "import random; print(random.randint(1, 2**32-1))")
OUTPUT_FILE="$OUTPUT_DIR/${FILENAME_PREFIX##*/}_${SEED}.jpg"

echo "CF AI REST API"
echo "Model:  $MODEL"
echo "Size:   ${WIDTH}x${HEIGHT}"
echo "Steps:  $STEPS"
echo "Seed:   $SEED"
echo "Generating..."

PAYLOAD=$(python3 -c "
import json, sys
print(json.dumps({
    'prompt': sys.argv[1],
    'steps':  int(sys.argv[2]),
    'seed':   int(sys.argv[3]),
    'width':  int(sys.argv[4]),
    'height': int(sys.argv[5]),
}))" "$PROMPT" "$STEPS" "$SEED" "$WIDTH" "$HEIGHT")

RAW_RESPONSE=$(mktemp)
HTTP_CODE=$(curl -s -w "%{http_code}" \
    -X POST "$API_URL" \
    -H "Authorization: Bearer $CF_API_TOKEN" \
    -H "Content-Type: application/json" \
    -d "$PAYLOAD" \
    -o "$RAW_RESPONSE")

if [ "$HTTP_CODE" != "200" ]; then
    echo "ERROR: CF AI API returned HTTP $HTTP_CODE"
    cat "$RAW_RESPONSE" 2>/dev/null || true
    rm -f "$RAW_RESPONSE"
    exit 1
fi

# Decode base64 image from JSON response
python3 -c "
import json, base64, sys
with open(sys.argv[1], 'rb') as f:
    data = json.load(f)
img = base64.b64decode(data['result']['image'])
with open(sys.argv[2], 'wb') as f:
    f.write(img)
" "$RAW_RESPONSE" "$OUTPUT_FILE"
rm -f "$RAW_RESPONSE"

# Blank image check
FILE_BYTES=$(wc -c < "$OUTPUT_FILE")
if [ "$FILE_BYTES" -lt 100000 ]; then
    echo "WARN: Output too small (${FILE_BYTES} bytes) — likely blank. Retrying with new seed..."
    rm -f "$OUTPUT_FILE"
    SEED=$(python3 -c "import random; print(random.randint(1, 2**32-1))")
    OUTPUT_FILE="$OUTPUT_DIR/${FILENAME_PREFIX##*/}_${SEED}.jpg"
    PAYLOAD=$(python3 -c "
import json, sys
print(json.dumps({
    'prompt': sys.argv[1],
    'steps':  int(sys.argv[2]),
    'seed':   int(sys.argv[3]),
    'width':  int(sys.argv[4]),
    'height': int(sys.argv[5]),
}))" "$PROMPT" "$STEPS" "$SEED" "$WIDTH" "$HEIGHT")
    RAW_RESPONSE=$(mktemp)
    HTTP_CODE=$(curl -s -w "%{http_code}" \
        -X POST "$API_URL" \
        -H "Authorization: Bearer $CF_API_TOKEN" \
        -H "Content-Type: application/json" \
        -d "$PAYLOAD" \
        -o "$RAW_RESPONSE")
    if [ "$HTTP_CODE" != "200" ]; then
        echo "ERROR: Retry returned HTTP $HTTP_CODE"
        rm -f "$RAW_RESPONSE"
        exit 1
    fi
    python3 -c "
import json, base64, sys
with open(sys.argv[1], 'rb') as f:
    data = json.load(f)
img = base64.b64decode(data['result']['image'])
with open(sys.argv[2], 'wb') as f:
    f.write(img)
" "$RAW_RESPONSE" "$OUTPUT_FILE"
    rm -f "$RAW_RESPONSE"
    FILE_BYTES=$(wc -c < "$OUTPUT_FILE")
    if [ "$FILE_BYTES" -lt 100000 ]; then
        echo "ERROR: Retry also blank. Giving up."
        rm -f "$OUTPUT_FILE"
        exit 1
    fi
    echo "Retry seed: $SEED"
fi

SIZE=$(du -h "$OUTPUT_FILE" | cut -f1)
echo "Done: $OUTPUT_FILE ($SIZE)"
echo "SEED:$SEED"