#!/bin/bash
# imggen — wrapper for sd-server image generation API
# Usage: imggen -p "your prompt" [options]

SERVER="http://localhost:8189"
PROMPT=""
STEPS=12
SEED=""
WIDTH=768
HEIGHT=512
OUTPUT=""
CFG=1.0
REF_IMAGE=""

usage() {
  cat <<USAGE
Usage: imggen -p "prompt" [options]

Options:
  -p, --prompt TEXT       Prompt (required)
  -s, --steps INT         Steps (default: 12)
  -S, --seed INT          Seed (default: random)
  -W, --width INT         Width (default: 768)
  -H, --height INT        Height (default: 512)
  -o, --output FILE       Output file (default: output-<seed>.png)
  -r, --ref IMAGE         Reference image for img2img
  -c, --cfg FLOAT         CFG scale (default: 1.0)
  --portrait              Shortcut for 512x768
  --landscape             Shortcut for 768x512 (default)
  --server URL            Server URL (default: http://localhost:8189)
  -h, --help              Show this help

Examples:
  imggen -p "a cat on a bicycle" --portrait -S 42
  imggen -p "photorealistic spy" -s 8 -o spy.png
  imggen -p "make this realistic" -r input.png
USAGE
  exit 0
}

die() { echo "Error: $1" >&2; exit 1; }

while [[ $# -gt 0 ]]; do
  case "$1" in
    -p|--prompt)    PROMPT="$2"; shift 2 ;;
    -s|--steps)     STEPS="$2"; shift 2 ;;
    -S|--seed)      SEED="$2"; shift 2 ;;
    -W|--width)     WIDTH="$2"; shift 2 ;;
    -H|--height)    HEIGHT="$2"; shift 2 ;;
    -o|--output)    OUTPUT="$2"; shift 2 ;;
    -r|--ref)       REF_IMAGE="$2"; shift 2 ;;
    -c|--cfg)       CFG="$2"; shift 2 ;;
    --portrait)     WIDTH=512; HEIGHT=768; shift ;;
    --landscape)    WIDTH=768; HEIGHT=512; shift ;;
    --server)       SERVER="$2"; shift 2 ;;
    -h|--help)      usage ;;
    *) die "Unknown option: $1" ;;
  esac
done

[[ -z "$PROMPT" ]] && die "Prompt is required (-p)"

if [[ -z "$SEED" ]]; then
  SEED=$RANDOM$RANDOM
fi

if [[ -z "$OUTPUT" ]]; then
  OUTPUT="output-${SEED}.png"
fi

curl -sf "$SERVER/" > /dev/null 2>&1 || die "sd-server not responding at $SERVER — is it running? (sudo systemctl start sd-server)"

echo "Generating: ${WIDTH}x${HEIGHT}, ${STEPS} steps, seed ${SEED}"
echo "Output: ${OUTPUT}"

if [[ -n "$REF_IMAGE" ]]; then
  [[ -f "$REF_IMAGE" ]] || die "Reference image not found: $REF_IMAGE"
  ENDPOINT="/sdapi/v1/img2img"
  REF_B64=$(base64 -w0 "$REF_IMAGE")
  PAYLOAD=$(printf '{"prompt":%s,"init_images":["%s"],"cfg_scale":%s,"steps":%d,"seed":%d,"width":%d,"height":%d}' \
    "$(echo "$PROMPT" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read().rstrip()))')" \
    "$REF_B64" "$CFG" "$STEPS" "$SEED" "$WIDTH" "$HEIGHT")
else
  ENDPOINT="/sdapi/v1/txt2img"
  PAYLOAD=$(printf '{"prompt":%s,"cfg_scale":%s,"steps":%d,"seed":%d,"width":%d,"height":%d}' \
    "$(echo "$PROMPT" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read().rstrip()))')" \
    "$CFG" "$STEPS" "$SEED" "$WIDTH" "$HEIGHT")
fi

START=$(date +%s)
RESPONSE=$(curl -sf -X POST "$SERVER$ENDPOINT" \
  -H "Content-Type: application/json" \
  -d "$PAYLOAD") || die "Request failed"

echo "$RESPONSE" | python3 -c "
import sys, json, base64
data = json.load(sys.stdin)
img = base64.b64decode(data['images'][0])
with open('$OUTPUT', 'wb') as f:
    f.write(img)
" || die "Failed to decode response"

END=$(date +%s)
SIZE=$(du -h "$OUTPUT" | cut -f1)
echo "Done in $((END - START))s — ${SIZE} → ${OUTPUT}"