#!/usr/bin/env bash
set -euo pipefail

# Load API key from .env
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

if [[ -f "$REPO_ROOT/.env" ]]; then
  # shellcheck disable=SC1091
  source "$REPO_ROOT/.env"
fi

usage() {
  echo "Usage: imagegen <prompt> [output.png]" >&2
  echo "" >&2
  echo "Generate an image from a text prompt using Gemini." >&2
  echo "Output defaults to 'output.png' if not specified." >&2
  exit 1
}

if [[ $# -lt 1 ]] || [[ "${1:-}" == "--help" ]] || [[ "${1:-}" == "-h" ]]; then
  usage
fi

if [[ -z "${GEMINI_API_KEY:-}" ]]; then
  echo "Error: GEMINI_API_KEY not set" >&2
  exit 1
fi

PROMPT="$1"
OUTPUT="${2:-output.png}"
MODEL="${IMAGEGEN_MODEL:-gemini-2.5-flash-image}"
API_URL="https://generativelanguage.googleapis.com/v1beta/models/${MODEL}:generateContent?key=${GEMINI_API_KEY}"

echo "Generating image for: ${PROMPT}" >&2
echo "Model: ${MODEL}" >&2

RESPONSE=$(curl -s --fail-with-body -X POST "$API_URL" \
  -H "Content-Type: application/json" \
  -d "$(jq -n --arg prompt "$PROMPT" '{
    contents: [{ parts: [{ text: $prompt }] }],
    generationConfig: { responseModalities: ["TEXT", "IMAGE"] }
  }')")

# Check for errors
ERROR=$(echo "$RESPONSE" | jq -r '.error.message // empty')
if [[ -n "$ERROR" ]]; then
  echo "API error: $ERROR" >&2
  exit 1
fi

# Extract base64 image data from the first image part
IMAGE_DATA=$(echo "$RESPONSE" | jq -r '
  .candidates[0].content.parts[]
  | select(.inlineData)
  | .inlineData.data
' | head -1)

if [[ -z "$IMAGE_DATA" ]]; then
  # Show any text response for debugging
  TEXT=$(echo "$RESPONSE" | jq -r '.candidates[0].content.parts[] | select(.text) | .text // empty')
  if [[ -n "$TEXT" ]]; then
    echo "Model response (no image): $TEXT" >&2
  else
    echo "No image in response. Raw:" >&2
    echo "$RESPONSE" | jq . >&2
  fi
  exit 1
fi

echo "$IMAGE_DATA" | base64 -d > "$OUTPUT"
echo "Saved to ${OUTPUT}" >&2