scripts/img_api_generate.sh

#!/bin/bash
#
# img_api_generate.sh - AI Image Generator via Pivoine API
# Generates images via Freepik API proxy and optionally swaps faces via FaceFusion
#
# Usage: ./img_api_generate.sh -p "prompt" [OPTIONS]
#

set -euo pipefail

# ============================================================================
# COLOR PALETTE - Creative Art Theme (Purple/Magenta)
# ============================================================================

RESET='\033[0m'

# Foreground Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
CYAN='\033[0;36m'

# Bold
BOLD_RED='\033[1;31m'
BOLD_GREEN='\033[1;32m'
BOLD_YELLOW='\033[1;33m'
BOLD_BLUE='\033[1;34m'
BOLD_MAGENTA='\033[1;35m'
BOLD_CYAN='\033[1;36m'
BOLD_WHITE='\033[1;37m'

# Styles
DIM='\033[2m'

# ============================================================================
# UNICODE CHARACTERS
# ============================================================================

CHECK_MARK="✓"
CROSS_MARK="✗"
PALETTE="🎨"
SPARKLES="✨"
CAMERA="📸"
FACE="🎭"
WARNING="⚠️"
INFO="ℹ️"
ARROW_RIGHT="→"
BOX_DOUBLE="═"

# ============================================================================
# CONSTANTS
# ============================================================================

BASE_URL="https://api.pivoine.art"
VALID_MODELS=("mystic" "flux-dev" "flux-pro" "seedream")
POLL_INTERVAL=3
POLL_TIMEOUT=600

# ============================================================================
# CONFIGURATION (defaults)
# ============================================================================

PROMPT=""
MODEL="flux-dev"
TOKEN=""
FACE_IMAGE=""
OUTPUT_FILE=""
NUM_IMAGES=1
SEED=""
NEGATIVE_PROMPT=""
GUIDANCE_SCALE=""
ASPECT_RATIO=""
INPUT_IMAGE=""
ASYNC_MODE=false
DRY_RUN=false
VERBOSE=false
UPSCALE=false

# Temp files for cleanup
_TMP_FILES=()

# ============================================================================
# LOGGING FUNCTIONS
# ============================================================================

print_banner() {
    local text="$1"
    local width=70
    local text_len=${#text}
    local padding=$(( (width - text_len) / 2 ))

    echo ""
    echo -e "${BOLD_MAGENTA}${BOX_DOUBLE}$(printf '%0.s═' $(seq 1 $width))${BOX_DOUBLE}${RESET}"
    echo -e "${BOLD_MAGENTA}║$(printf '%*s' $padding '')${BOLD_CYAN}${text}$(printf '%*s' $((width - padding - text_len)) '')${BOLD_MAGENTA}║${RESET}"
    echo -e "${BOLD_MAGENTA}${BOX_DOUBLE}$(printf '%0.s═' $(seq 1 $width))${BOX_DOUBLE}${RESET}"
    echo ""
}

print_section() {
    local text="$1"
    echo -e "\n${BOLD_MAGENTA}» ${text}${RESET}"
    echo -e "${MAGENTA}$(printf '%0.s─' $(seq 1 70))${RESET}"
}

print_success() {
    echo -e "${BOLD_GREEN}${CHECK_MARK} $1${RESET}"
}

print_error() {
    echo -e "${BOLD_RED}${CROSS_MARK} $1${RESET}" >&2
}

print_warning() {
    echo -e "${BOLD_YELLOW}${WARNING} $1${RESET}"
}

print_info() {
    echo -e "${BOLD_CYAN}${INFO} $1${RESET}"
}

print_detail() {
    echo -e "  ${DIM}${MAGENTA}${ARROW_RIGHT} $1${RESET}"
}

print_field() {
    local icon="$1"
    local label="$2"
    local value="$3"
    printf "  ${icon} ${BOLD_WHITE}%-15s${RESET} %s\n" "${label}:" "$value"
}

print_verbose() {
    if [[ "$VERBOSE" == true ]]; then
        echo -e "  ${DIM}[verbose] $1${RESET}" >&2
    fi
}

# ============================================================================
# SPINNER
# ============================================================================

spinner() {
    local pid=$1
    local message="${2:-Processing}"
    local spinchars='⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'
    local i=0

    tput civis 2>/dev/null || true
    while kill -0 "$pid" 2>/dev/null; do
        printf "\r  ${BOLD_MAGENTA}${spinchars:i++%${#spinchars}:1}${RESET} ${BOLD_WHITE}%s...${RESET}" "$message"
        sleep 0.1
    done
    tput cnorm 2>/dev/null || true
    printf "\r%*s\r" 60 ""
}

# ============================================================================
# CLEANUP
# ============================================================================

cleanup() {
    tput cnorm 2>/dev/null || true
    for f in "${_TMP_FILES[@]}"; do
        rm -f "$f" 2>/dev/null || true
    done
}

trap cleanup EXIT

# ============================================================================
# DEPENDENCY CHECK
# ============================================================================

check_dependencies() {
    local missing=()

    for cmd in curl jq base64; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done

    # file is only needed for --image (mime type detection)
    if [[ -n "$INPUT_IMAGE" ]] && ! command -v file &>/dev/null; then
        missing+=("file")
    fi

    if [[ ${#missing[@]} -gt 0 ]]; then
        print_error "Missing required dependencies: ${missing[*]}"
        echo ""
        echo "Install with:"
        echo -e "  ${DIM}sudo apt install ${missing[*]}${RESET}"
        exit 1
    fi
}

# ============================================================================
# TOKEN RESOLUTION
# ============================================================================

resolve_token() {
    # 1. Already set via --token
    if [[ -n "$TOKEN" ]]; then
        print_verbose "Token from --token argument"
        return 0
    fi

    # 2. Environment variable
    if [[ -n "${PIVOINE_API_TOKEN:-}" ]]; then
        TOKEN="$PIVOINE_API_TOKEN"
        print_verbose "Token from PIVOINE_API_TOKEN env var"
        return 0
    fi

    # 3. Config file
    local token_file="${HOME}/.config/pivoine/token"
    if [[ -f "$token_file" ]]; then
        TOKEN="$(tr -d '[:space:]' < "$token_file")"
        if [[ -n "$TOKEN" ]]; then
            print_verbose "Token from ${token_file}"
            return 0
        fi
    fi

    print_error "API token not found"
    print_info "Provide via --token, PIVOINE_API_TOKEN env var, or ~/.config/pivoine/token"
    exit 1
}

# ============================================================================
# INPUT VALIDATION
# ============================================================================

validate_inputs() {
    local errors=0

    # Prompt is required
    if [[ -z "$PROMPT" ]]; then
        print_error "Prompt is required (-p, --prompt)"
        ((errors++))
    elif [[ ${#PROMPT} -gt 4000 ]]; then
        print_error "Prompt too long (${#PROMPT} chars, max 4000)"
        ((errors++))
    fi

    # Model must be valid
    local model_valid=false
    for m in "${VALID_MODELS[@]}"; do
        if [[ "$MODEL" == "$m" ]]; then
            model_valid=true
            break
        fi
    done
    if [[ "$model_valid" == false ]]; then
        print_error "Invalid model: $MODEL (valid: ${VALID_MODELS[*]})"
        ((errors++))
    fi

    # Face image must exist if specified
    if [[ -n "$FACE_IMAGE" && ! -f "$FACE_IMAGE" ]]; then
        print_error "Face image not found: $FACE_IMAGE"
        ((errors++))
    fi

    # Input image must exist if specified
    if [[ -n "$INPUT_IMAGE" && ! -f "$INPUT_IMAGE" ]]; then
        print_error "Input image not found: $INPUT_IMAGE"
        ((errors++))
    fi

    # Numeric range checks
    if [[ -n "$NUM_IMAGES" ]]; then
        if ! [[ "$NUM_IMAGES" =~ ^[1-4]$ ]]; then
            print_error "Number of images must be 1-4 (got: $NUM_IMAGES)"
            ((errors++))
        fi
    fi

    if [[ -n "$GUIDANCE_SCALE" ]]; then
        if ! awk "BEGIN {exit ($GUIDANCE_SCALE >= 1.0 && $GUIDANCE_SCALE <= 20.0) ? 0 : 1}" 2>/dev/null; then
            print_error "Guidance scale must be 1.0-20.0 (got: $GUIDANCE_SCALE)"
            ((errors++))
        fi
        if [[ "$MODEL" != "flux-dev" && "$MODEL" != "flux-pro" ]]; then
            print_warning "Guidance scale is only supported for flux models (ignoring)"
            GUIDANCE_SCALE=""
        fi
    fi

    if [[ -n "$NEGATIVE_PROMPT" && "$MODEL" != "mystic" ]]; then
        print_warning "Negative prompt is only supported for mystic model (ignoring)"
        NEGATIVE_PROMPT=""
    fi

    if [[ -n "$ASPECT_RATIO" && "$MODEL" != "seedream" ]]; then
        print_warning "Aspect ratio is only supported for seedream model (ignoring)"
        ASPECT_RATIO=""
    fi

    if ((errors > 0)); then
        echo ""
        print_error "Found $errors error(s). Aborting."
        exit 1
    fi
}

# ============================================================================
# OUTPUT FILENAME GENERATION
# ============================================================================

generate_output_filename() {
    local suffix="${1:-}"
    local timestamp
    timestamp=$(date +%Y%m%d_%H%M%S)
    local seed_part="${SEED:-0}"
    local base="img_${MODEL}_${timestamp}_${seed_part}"

    if [[ -n "$suffix" ]]; then
        echo "${base}_${suffix}.png"
    else
        echo "${base}.png"
    fi
}

# ============================================================================
# PAYLOAD BUILDER
# ============================================================================

build_payload() {
    local payload

    # Start with required fields
    payload=$(jq -n \
        --arg prompt "$PROMPT" \
        --argjson num_images "$NUM_IMAGES" \
        '{prompt: $prompt, num_images: $num_images}')

    # Optional: seed
    if [[ -n "$SEED" ]]; then
        payload=$(echo "$payload" | jq --argjson seed "$SEED" '. + {seed: $seed}')
    fi

    # Optional: negative prompt (mystic only)
    if [[ -n "$NEGATIVE_PROMPT" ]]; then
        payload=$(echo "$payload" | jq --arg neg "$NEGATIVE_PROMPT" '. + {negative_prompt: $neg}')
    fi

    # Optional: guidance scale (flux models)
    if [[ -n "$GUIDANCE_SCALE" ]]; then
        payload=$(echo "$payload" | jq --argjson gs "$GUIDANCE_SCALE" '. + {guidance_scale: $gs}')
    fi

    # Optional: aspect ratio (seedream)
    if [[ -n "$ASPECT_RATIO" ]]; then
        payload=$(echo "$payload" | jq --arg ar "$ASPECT_RATIO" '. + {aspect_ratio: $ar}')
    fi

    # Optional: input image (img2img)
    if [[ -n "$INPUT_IMAGE" ]]; then
        local b64
        b64=$(base64 -w 0 "$INPUT_IMAGE")
        local mime
        mime=$(file -b --mime-type "$INPUT_IMAGE")
        payload=$(echo "$payload" | jq \
            --arg img "data:${mime};base64,${b64}" \
            '. + {image: $img}')
    fi

    echo "$payload"
}

# ============================================================================
# API HELPERS
# ============================================================================

api_curl() {
    local method="$1"
    local url="$2"
    shift 2
    local extra_args=("$@")

    local tmp_body tmp_headers
    tmp_body=$(mktemp)
    tmp_headers=$(mktemp)
    _TMP_FILES+=("$tmp_body" "$tmp_headers")

    local http_code
    http_code=$(curl -s -w "%{http_code}" \
        -X "$method" \
        -H "X-Api-Key: ${TOKEN}" \
        -D "$tmp_headers" \
        -o "$tmp_body" \
        "${extra_args[@]}" \
        "$url") || {
        print_error "curl failed (network error)"
        return 1
    }

    print_verbose "HTTP ${http_code} from ${method} ${url}"

    # Handle common HTTP errors
    case "$http_code" in
        2[0-9][0-9])
            # Success
            cat "$tmp_body"
            return 0
            ;;
        401)
            print_error "Authentication failed (HTTP 401)"
            print_info "Check your API token (--token, PIVOINE_API_TOKEN, or ~/.config/pivoine/token)"
            return 1
            ;;
        429)
            print_warning "Rate limited (HTTP 429). Retrying in 2s..."
            sleep 2
            http_code=$(curl -s -w "%{http_code}" \
                -X "$method" \
                -H "X-Api-Key: ${TOKEN}" \
                -D "$tmp_headers" \
                -o "$tmp_body" \
                "${extra_args[@]}" \
                "$url") || {
                print_error "curl failed on retry (network error)"
                return 1
            }
            if [[ "$http_code" =~ ^2 ]]; then
                cat "$tmp_body"
                return 0
            fi
            print_error "Still rate limited after retry (HTTP ${http_code})"
            cat "$tmp_body" >&2
            return 1
            ;;
        *)
            print_error "API error (HTTP ${http_code})"
            local err_body
            err_body=$(cat "$tmp_body" 2>/dev/null || echo "(empty)")
            if echo "$err_body" | jq . &>/dev/null; then
                echo "$err_body" | jq -r '.detail // .message // .error // .' >&2
            else
                echo "  $err_body" >&2
            fi
            return 1
            ;;
    esac
}

content_type_to_ext() {
    local ct="$1"
    case "$ct" in
        *image/jpeg*)  echo "jpg" ;;
        *image/png*)   echo "png" ;;
        *image/webp*)  echo "webp" ;;
        *image/gif*)   echo "gif" ;;
        *)             echo "" ;;
    esac
}

# Global set by api_curl_binary after each call
_RESPONSE_CONTENT_TYPE=""

api_curl_binary() {
    local method="$1"
    local url="$2"
    local output_path="$3"
    shift 3
    local extra_args=("$@")

    local tmp_headers
    tmp_headers=$(mktemp)
    _TMP_FILES+=("$tmp_headers")

    _RESPONSE_CONTENT_TYPE=""

    local http_code
    http_code=$(curl -s -w "%{http_code}" \
        -X "$method" \
        -H "X-Api-Key: ${TOKEN}" \
        -D "$tmp_headers" \
        -o "$output_path" \
        "${extra_args[@]}" \
        "$url") || {
        print_error "curl failed (network error)"
        return 1
    }

    print_verbose "HTTP ${http_code} from ${method} ${url}"

    # Extract content-type from response headers
    _RESPONSE_CONTENT_TYPE=$(grep -i '^content-type:' "$tmp_headers" 2>/dev/null | tail -1 | cut -d: -f2- | tr -d '[:space:]')
    print_verbose "Content-Type: ${_RESPONSE_CONTENT_TYPE}"

    if [[ "$http_code" =~ ^2 ]]; then
        return 0
    fi

    print_error "Download failed (HTTP ${http_code})"
    # If the response was a JSON error, show it
    if head -c 1 "$output_path" 2>/dev/null | grep -q '{'; then
        cat "$output_path" >&2
        rm -f "$output_path"
    fi
    return 1
}

# Rename file to match the actual content-type extension.
# Updates the path variable named by $2 (passed by reference).
fix_extension() {
    local current_path="$1"
    local -n path_ref="$2"
    local detected_ext
    detected_ext=$(content_type_to_ext "$_RESPONSE_CONTENT_TYPE")

    if [[ -z "$detected_ext" ]]; then
        return 0
    fi

    local current_ext="${current_path##*.}"
    if [[ "$current_ext" == "$detected_ext" ]]; then
        return 0
    fi

    local new_path="${current_path%.*}.${detected_ext}"
    mv "$current_path" "$new_path"
    path_ref="$new_path"
    print_verbose "Renamed to match content-type: ${new_path}"
}

# ============================================================================
# IMAGE GENERATION
# ============================================================================

generate_image() {
    local payload="$1"
    local url="${BASE_URL}/freepik/generate/image/${MODEL}"

    if [[ "$ASYNC_MODE" == false ]]; then
        url="${url}?sync=true"
    fi

    if [[ "$DRY_RUN" == true ]]; then
        print_section "Dry Run - Generate Image"
        echo ""
        echo -e "${DIM}curl -s -X POST \\"
        echo "  -H 'X-Api-Key: \${TOKEN}' \\"
        echo "  -H 'Content-Type: application/json' \\"
        echo "  -d '$(echo "$payload" | jq -c .)' \\"
        echo -e "  '${url}'${RESET}"
        echo ""
        echo -e "${DIM}curl -s -X GET \\"
        echo "  -H 'X-Api-Key: \${TOKEN}' \\"
        echo "  -o '\${OUTPUT_FILE}' \\"
        echo -e "  '${BASE_URL}/freepik/tasks/\${TASK_ID}/result'${RESET}"
        echo ""
        return 0
    fi

    print_section "Generating Image"
    print_field "$PALETTE" "Model" "$MODEL"
    print_field "$SPARKLES" "Prompt" "$(echo "$PROMPT" | head -c 80)$([ ${#PROMPT} -gt 80 ] && echo '...')"
    [[ -n "$SEED" ]] && print_field " " "Seed" "$SEED"
    [[ -n "$NEGATIVE_PROMPT" ]] && print_field " " "Negative" "$(echo "$NEGATIVE_PROMPT" | head -c 60)..."
    [[ -n "$GUIDANCE_SCALE" ]] && print_field " " "Guidance" "$GUIDANCE_SCALE"
    [[ -n "$ASPECT_RATIO" ]] && print_field " " "Aspect" "$ASPECT_RATIO"
    [[ -n "$INPUT_IMAGE" ]] && print_field " " "Input Image" "$INPUT_IMAGE"
    echo ""

    local response
    if [[ "$ASYNC_MODE" == false ]]; then
        # Synchronous mode - wait with spinner
        local tmp_resp
        tmp_resp=$(mktemp)
        _TMP_FILES+=("$tmp_resp")

        (api_curl POST "$url" \
            -H "Content-Type: application/json" \
            -d "$payload" > "$tmp_resp" 2>&1) &
        local pid=$!
        spinner "$pid" "Generating image (sync)"
        wait "$pid" || {
            cat "$tmp_resp" >&2
            return 1
        }
        response=$(cat "$tmp_resp")
    else
        # Async mode - post and poll
        response=$(api_curl POST "$url" \
            -H "Content-Type: application/json" \
            -d "$payload") || return 1
    fi

    # Parse task_id from response
    TASK_ID=$(echo "$response" | jq -r '.task_id // empty' 2>/dev/null) || {
        print_error "Failed to parse response"
        print_verbose "Raw response: $response"
        return 1
    }

    if [[ -z "$TASK_ID" ]]; then
        print_error "No task_id in response"
        echo "$response" | jq . 2>/dev/null || echo "$response" >&2
        return 1
    fi

    print_success "Task ID: ${TASK_ID}"

    # If async, poll for completion
    if [[ "$ASYNC_MODE" == true ]]; then
        poll_task "$TASK_ID" || return 1
    fi

    # Check status from sync response
    local status
    status=$(echo "$response" | jq -r '.status // "unknown"' 2>/dev/null)
    if [[ "$status" != "completed" && "$ASYNC_MODE" == false ]]; then
        # Might still need polling even in sync mode if server returned early
        print_info "Status: ${status} - polling for completion..."
        poll_task "$TASK_ID" || return 1
    fi
}

# ============================================================================
# TASK POLLING (async mode)
# ============================================================================

poll_task() {
    local task_id="$1"
    local url="${BASE_URL}/freepik/tasks/${task_id}"
    local start_time elapsed status
    start_time=$(date +%s)

    local spinchars='⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'
    local i=0

    tput civis 2>/dev/null || true

    while true; do
        elapsed=$(( $(date +%s) - start_time ))
        if [[ $elapsed -ge $POLL_TIMEOUT ]]; then
            tput cnorm 2>/dev/null || true
            printf "\r%*s\r" 60 ""
            print_error "Polling timed out after ${POLL_TIMEOUT}s"
            return 1
        fi

        local response
        response=$(curl -s \
            -H "X-Api-Key: ${TOKEN}" \
            "$url") || {
            printf "\r%*s\r" 60 ""
            print_warning "Poll request failed, retrying..."
            sleep "$POLL_INTERVAL"
            continue
        }

        status=$(echo "$response" | jq -r '.status // "unknown"' 2>/dev/null)

        case "$status" in
            completed)
                tput cnorm 2>/dev/null || true
                printf "\r%*s\r" 60 ""
                print_success "Task completed (${elapsed}s)"
                return 0
                ;;
            failed|error)
                tput cnorm 2>/dev/null || true
                printf "\r%*s\r" 60 ""
                local detail
                detail=$(echo "$response" | jq -r '.detail // .error // "unknown error"' 2>/dev/null)
                print_error "Task failed: $detail"
                return 1
                ;;
            *)
                local sc="${spinchars:i++%${#spinchars}:1}"
                printf "\r  %b%s%b %bWaiting for task... %b(%ss, status: %s)%b" \
                    "$BOLD_MAGENTA" "$sc" "$RESET" "$BOLD_WHITE" "$DIM" "$elapsed" "$status" "$RESET"
                sleep "$POLL_INTERVAL"
                ;;
        esac
    done
}

# ============================================================================
# DOWNLOAD RESULT
# ============================================================================

download_result() {
    local task_id="$1"
    local -n _dl_output_path="$2"
    local url="${BASE_URL}/freepik/tasks/${task_id}/result"

    if [[ "$DRY_RUN" == true ]]; then
        echo ""
        echo -e "${DIM}curl -s -X GET \\"
        echo "  -H 'X-Api-Key: \${TOKEN}' \\"
        echo "  -o '${_dl_output_path}' \\"
        echo -e "  '${url}'${RESET}"
        echo ""
        return 0
    fi

    print_detail "Downloading result..."

    api_curl_binary GET "$url" "$_dl_output_path" || return 1

    if [[ -f "$_dl_output_path" && -s "$_dl_output_path" ]]; then
        fix_extension "$_dl_output_path" _dl_output_path
        local size
        size=$(du -h "$_dl_output_path" | cut -f1)
        print_success "Saved: ${BOLD_WHITE}${_dl_output_path}${RESET} (${size})"
        return 0
    else
        print_error "Downloaded file is empty or missing"
        rm -f "$_dl_output_path" 2>/dev/null || true
        return 1
    fi
}

# ============================================================================
# FACE SWAP
# ============================================================================

swap_face() {
    local target_image="$1"
    local source_face="$2"
    local -n _fs_output_path="$3"
    local url="${BASE_URL}/facefusion/process"

    if [[ "$DRY_RUN" == true ]]; then
        print_section "Dry Run - Face Swap curl command"
        echo ""
        echo -e "${DIM}curl -s -X POST \\"
        echo "  -H 'X-Api-Key: \${TOKEN}' \\"
        echo "  -F 'target=@${target_image}' \\"
        echo "  -F 'source=@${source_face}' \\"
        echo "  -F 'options={\"processors\":[\"face_swapper\",\"face_enhancer\"],\"face_swapper\":{\"model\":\"hyperswap_1a_256\"},\"face_enhancer\":{\"model\":\"gfpgan_1.4\",\"blend\":80}}' \\"
        echo -e "  '${url}'${RESET}"
        echo ""
        return 0
    fi

    print_section "Face Swap"
    print_field "$FACE" "Source Face" "$source_face"
    print_field "$CAMERA" "Target" "$target_image"
    echo ""

    local options='{"processors":["face_swapper","face_enhancer"],"face_swapper":{"model":"hyperswap_1a_256"},"face_enhancer":{"model":"gfpgan_1.4","blend":80}}'

    local tmp_resp
    tmp_resp=$(mktemp)
    _TMP_FILES+=("$tmp_resp")

    (api_curl_binary POST "$url" "$_fs_output_path" \
        -F "target=@${target_image}" \
        -F "source=@${source_face}" \
        -F "options=${options}" > "$tmp_resp" 2>&1) &
    local pid=$!
    spinner "$pid" "Swapping face"
    wait "$pid" || {
        cat "$tmp_resp" >&2
        print_warning "Face swap failed, keeping original image"
        return 1
    }

    if [[ -f "$_fs_output_path" && -s "$_fs_output_path" ]]; then
        fix_extension "$_fs_output_path" _fs_output_path
        local size
        size=$(du -h "$_fs_output_path" | cut -f1)
        print_success "Face swap saved: ${BOLD_WHITE}${_fs_output_path}${RESET} (${size})"
        return 0
    else
        print_warning "Face swap result is empty, keeping original image"
        rm -f "$_fs_output_path" 2>/dev/null || true
        return 1
    fi
}

# ============================================================================
# Real-ESRGAN UPSCALE
# ============================================================================
realesrgan_upscale() {
    local input_image="$1"
    local -n _up_output_path="$2"
    local url="${BASE_URL}/realesrgan/upscale"

    if [[ "$DRY_RUN" == true ]]; then
        print_section "Dry Run - Real-ESRGAN Upscale curl command"
        echo ""
        echo -e "${DIM}curl -s -X POST \\\"\n"
        echo "  -H 'X-Api-Key: \\${TOKEN}' \\\"\n"
        echo "  -F 'image=@${input_image}' \\\"\n"
        echo "  -F 'model=RealESRGAN_x4plus' \\\"\n"
        echo -e "  '${url}'${RESET}"
        echo ""
        return 0
    fi

    print_section "Upscaling (Real-ESRGAN)"
    print_field "$PALETTE" "Model" "RealESRGAN_x4plus"
    print_field "$CAMERA" "Input" "$input_image"
    echo ""

    local tmp_resp
    tmp_resp=$(mktemp)
    _TMP_FILES+=("$tmp_resp")

    (api_curl_binary POST "$url" "${_up_output_path}" \
        -F "image=@${input_image}" \
        -F "model=RealESRGAN_x4plus" > "$tmp_resp" 2>&1) &
    local pid=$!
    spinner "$pid" "Upscaling image"
    wait "$pid" || {
        cat "$tmp_resp" >&2
        print_error "Upscale failed"
        return 1
    }

    if [[ -f "${_up_output_path}" && -s "${_up_output_path}" ]]; then
        fix_extension "${_up_output_path}" _up_output_path
        local size
        size=$(du -h "${_up_output_path}" | cut -f1)
        print_success "Upscaled saved: ${BOLD_WHITE}${_up_output_path}${RESET} (${size})"
        return 0
    else
        print_error "Upscale result is empty or missing"
        rm -f "${_up_output_path}" 2>/dev/null || true
        return 1
    fi
}

# ============================================================================
# HELP
# ============================================================================

show_help() {
    cat << 'EOF'
img_api_generate.sh - AI Image Generator via Pivoine API

Usage: img_api_generate.sh -p "prompt" [OPTIONS]

REQUIRED:
  -p, --prompt TEXT           Text prompt for image generation (1-4000 chars)

MODEL & GENERATION:
  -m, --model MODEL           Model name (default: flux-dev)
                              Available: mystic, flux-dev, flux-pro, seedream
  -n, --num-images N          Number of images 1-4 (default: 1)
  -s, --seed N                Seed for reproducibility
      --negative-prompt TEXT   Negative prompt (mystic model only)
      --guidance-scale N       Guidance scale 1.0-20.0 (flux models only)
      --aspect-ratio RATIO     Aspect ratio (seedream model only)
      --image FILE             Input image for img2img (base64-encoded automatically)

AUTHENTICATION:
  -t, --token TOKEN           API token (X-Api-Key)
                              Falls back to PIVOINE_API_TOKEN env var,
                              then ~/.config/pivoine/token file

FACE SWAP:
  -f, --face FILE             Source face image - triggers face swap on result

OUTPUT:
  -o, --output FILE           Output file path (default: auto-generated)
    -u, --upscale               Upscale final image using Real-ESRGAN (RealESRGAN_x4plus)

MODES:
      --async                 Use async mode with polling instead of sync
      --dry-run               Show curl commands without executing
      --verbose               Verbose output
  -h, --help                  Show this help message

EXAMPLES:
  # Generate with flux-dev (default model)
  ./img_api_generate.sh -p "a cat sitting on a rainbow"

  # Generate with mystic model and negative prompt
  ./img_api_generate.sh -p "photorealistic portrait" -m mystic \
      --negative-prompt "blurry, low quality"

  # Generate and face swap
  ./img_api_generate.sh -p "portrait of a person" -m mystic -f face.jpg

  # Dry run to preview API calls
  ./img_api_generate.sh --dry-run -p "a cat" -m flux-dev

  # Use a specific seed for reproducibility
  ./img_api_generate.sh -p "landscape at sunset" -s 42 -m flux-pro

EOF
}

# ============================================================================
# ARGUMENT PARSING
# ============================================================================

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -p|--prompt)
                PROMPT="$2"
                shift 2
                ;;
            -m|--model)
                MODEL="$2"
                shift 2
                ;;
            -t|--token)
                TOKEN="$2"
                shift 2
                ;;
            -f|--face)
                FACE_IMAGE="$2"
                shift 2
                ;;
            -o|--output)
                OUTPUT_FILE="$2"
                shift 2
                ;;
            -n|--num-images)
                NUM_IMAGES="$2"
                shift 2
                ;;
            -s|--seed)
                SEED="$2"
                shift 2
                ;;
            --negative-prompt)
                NEGATIVE_PROMPT="$2"
                shift 2
                ;;
            --guidance-scale)
                GUIDANCE_SCALE="$2"
                shift 2
                ;;
            --aspect-ratio)
                ASPECT_RATIO="$2"
                shift 2
                ;;
            --image)
                INPUT_IMAGE="$2"
                shift 2
                ;;
            --async)
                ASYNC_MODE=true
                shift
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            -u|--upscale)
                UPSCALE=true
                shift
                ;;
            --verbose)
                VERBOSE=true
                shift
                ;;
            -h|--help)
                show_help
                exit 0
                ;;
            -*)
                print_error "Unknown option: $1"
                echo "Use --help for usage information"
                exit 1
                ;;
            *)
                print_error "Unexpected argument: $1"
                echo "Use --help for usage information"
                exit 1
                ;;
        esac
    done
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    parse_args "$@"

    # Show help if no prompt
    if [[ -z "$PROMPT" ]]; then
        show_help
        exit 0
    fi

    # Check dependencies
    check_dependencies

    # Resolve API token
    resolve_token

    # Validate inputs
    validate_inputs

    # Banner
    print_banner "${PALETTE} Pivoine Image Generator ${SPARKLES}"

    if [[ "$DRY_RUN" == true ]]; then
        echo -e "${BOLD_YELLOW}${WARNING}  DRY-RUN MODE - No API calls will be made ${WARNING}${RESET}\n"
    fi

    # Build payload
    local payload
    payload=$(build_payload)
    print_verbose "Payload: $(echo "$payload" | jq -c .)"

    # Generate image (sets global TASK_ID)
    TASK_ID=""
    generate_image "$payload" || exit 1

    if [[ "$DRY_RUN" == true ]]; then
        # Show face swap dry-run if applicable
        local dry_input="/tmp/generated.png"
        if [[ -n "$FACE_IMAGE" ]]; then
            tmp_faceswap="/tmp/faceswap.png"
            swap_face "/tmp/generated.png" "$FACE_IMAGE" tmp_faceswap
            dry_input="$tmp_faceswap"
        fi

        if [[ "$UPSCALE" == true ]]; then
            tmp_upscaled="/tmp/upscaled.png"
            realesrgan_upscale "$dry_input" tmp_upscaled
        fi

        print_banner "${SPARKLES} Dry Run Complete ${SPARKLES}"
        exit 0
    fi

    # Determine output filename
    local output_path
    if [[ -n "$OUTPUT_FILE" ]]; then
        output_path="$OUTPUT_FILE"
    else
        output_path="$(generate_output_filename)"
    fi

    # Download result (output_path may be renamed to match content-type)
    download_result "$TASK_ID" output_path || exit 1

    # Face swap if requested
    if [[ -n "$FACE_IMAGE" ]]; then
        local faceswap_path
        if [[ -n "$OUTPUT_FILE" ]]; then
            # Insert _faceswap before extension
            local ext="${OUTPUT_FILE##*.}"
            local base="${OUTPUT_FILE%.*}"
            faceswap_path="${base}_faceswap.${ext}"
        else
            faceswap_path="$(generate_output_filename "faceswap")"
        fi

        swap_face "$output_path" "$FACE_IMAGE" faceswap_path || true
    fi

    # Determine final input for optional upscale (prefer faceswap if present)
    local final_input="$output_path"
    if [[ -n "${faceswap_path:-}" && -f "${faceswap_path}" ]]; then
        final_input="${faceswap_path}"
    fi

    # Upscale with Real-ESRGAN if requested
    if [[ "$UPSCALE" == true ]]; then
        local upscale_path
        if [[ -n "$OUTPUT_FILE" ]]; then
            local ext="${OUTPUT_FILE##*.}"
            local base="${OUTPUT_FILE%.*}"
            upscale_path="${base}_upscaled.${ext}"
        else
            upscale_path="$(generate_output_filename "upscaled")"
        fi

        realesrgan_upscale "$final_input" upscale_path || exit 1

        # Use the upscaled image for summary/output
        output_path="$upscale_path"
    fi

    # Summary
    print_section "Summary"
    print_field "$PALETTE" "Model" "$MODEL"
    print_field "$CAMERA" "Output" "$output_path"
    if [[ -n "$FACE_IMAGE" && -f "${faceswap_path:-}" ]]; then
        print_field "$FACE" "Face Swap" "$faceswap_path"
    fi
    [[ -n "$SEED" ]] && print_field " " "Seed" "$SEED"
    echo ""

    print_banner "${SPARKLES} Complete ${SPARKLES}"
}

main "$@"