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

usage() {
  cat >&2 <<'USAGE'
Usage:
  async-media.sh image MODEL PROMPT OUTPUT [SIZE]
  async-media.sh video MODEL PROMPT OUTPUT [SECONDS] [SIZE]

Examples:
  export SOURCESDATA_BASE_URL='<YOUR_SOURCESDATA_BASE_URL>'
  export SOURCESDATA_API_KEY='<YOUR_SOURCESDATA_API_KEY>'
  ./examples/bash/async-media.sh image gpt-image-2 \
    'A clean studio photograph of a glass perfume bottle' output.png 1024x1024
  ./examples/bash/async-media.sh video seedance-2.0-fast \
    'A paper boat drifting through a rain puddle' output.mp4 4 1280x720
USAGE
  exit 2
}

require_command() {
  command -v "$1" >/dev/null 2>&1 || {
    printf 'Required command not found: %s\n' "$1" >&2
    exit 2
  }
}

for command_name in curl jq base64; do
  require_command "$command_name"
done

: "${SOURCESDATA_BASE_URL:?Set SOURCESDATA_BASE_URL='<YOUR_SOURCESDATA_BASE_URL>'}"
: "${SOURCESDATA_API_KEY:?Set SOURCESDATA_API_KEY='<YOUR_SOURCESDATA_API_KEY>'}"

[[ $# -ge 4 ]] || usage
media_type=$1
model=$2
prompt=$3
output=$4
base_url=${SOURCESDATA_BASE_URL%/}
submit_path=/v1/videos

case "$media_type" in
  image)
    size=${5:-1024x1024}
    submit_path=/v1/images/generations
    payload=$(jq -n \
      --arg model "$model" \
      --arg prompt "$prompt" \
      --arg size "$size" \
      '{model: $model, prompt: $prompt, size: $size, n: 1, response_format: "url", async: true}')
    ;;
  video)
    seconds=${5:-4}
    size=${6:-1280x720}
    [[ $seconds =~ ^[0-9]+$ ]] || {
      echo "SECONDS must be a positive integer" >&2
      exit 2
    }
    payload=$(jq -n \
      --arg model "$model" \
      --arg prompt "$prompt" \
      --arg size "$size" \
      --argjson seconds "$seconds" \
      '{model: $model, prompt: $prompt, seconds: $seconds, size: $size}')
    ;;
  *) usage ;;
esac

tmp_dir=$(mktemp -d)
trap 'rm -rf "$tmp_dir"' EXIT

submit_body=$tmp_dir/submit.json
submit_headers=$tmp_dir/submit.headers
poll_body=$tmp_dir/poll.json
poll_headers=$tmp_dir/poll.headers

curl --fail-with-body --silent --show-error \
  --connect-timeout 15 --max-time 120 \
  -D "$submit_headers" -o "$submit_body" \
  "$base_url$submit_path" \
  -H "Authorization: Bearer $SOURCESDATA_API_KEY" \
  -H 'Content-Type: application/json' \
  --data "$payload"

task_id=$(jq -r '.id // .task_id // .data.id // .data.task_id // empty' "$submit_body")
if [[ ! $task_id =~ ^task_[A-Za-z0-9_-]+$ ]]; then
  echo "The response did not contain a valid SourcesData task ID" >&2
  jq '{status, error, code}' "$submit_body" >&2 || true
  exit 1
fi

printf 'Task: %s\n' "$task_id"

retry_after() {
  local headers=$1
  local value
  value=$(awk 'BEGIN { IGNORECASE=1 } /^Retry-After:/ { gsub("\\r", "", $2); print $2 }' "$headers" | tail -n 1)
  if [[ $value =~ ^[0-9]+$ ]] && (( value >= 1 && value <= 60 )); then
    printf '%s' "$value"
  else
    printf '4'
  fi
}

initial_wait=$(retry_after "$submit_headers")
printf 'First poll in %s seconds\n' "$initial_wait"
sleep "$initial_wait"

decode_base64() {
  if base64 --decode >/dev/null 2>&1 <<<''; then
    base64 --decode
  else
    base64 -D
  fi
}

download_image_result() {
  local body=$1
  local image_url image_base64
  image_url=$(jq -r '.url // .data[0].url // .data.url // .data.data[0].url // empty' "$body")
  image_base64=$(jq -r '.b64_json // .data[0].b64_json // .data.b64_json // .data.data[0].b64_json // empty' "$body")

  if [[ -n $image_url ]]; then
    case "$image_url" in
      "$base_url"/*) ;;
      *)
        echo "Refusing to download a non-SourcesData image URL" >&2
        return 1
        ;;
    esac
    curl --fail-with-body --silent --show-error --location \
      --connect-timeout 15 --max-time 600 \
      "$image_url" --output "$output"
    return 0
  fi

  if [[ -n $image_base64 ]]; then
    printf '%s' "$image_base64" | decode_base64 >"$output"
    return 0
  fi

  return 1
}

for ((attempt = 1; attempt <= 225; attempt += 1)); do
  : >"$poll_headers"
  curl --fail-with-body --silent --show-error \
    --connect-timeout 15 --max-time 60 \
    -D "$poll_headers" -o "$poll_body" \
    "$base_url/v1/videos/$task_id" \
    -H "Authorization: Bearer $SOURCESDATA_API_KEY"

  status=$(jq -r '.status // .data.status // empty | ascii_downcase' "$poll_body")
  progress=$(jq -r '.progress // .data.progress // empty' "$poll_body")
  printf 'Poll %d: %s%s\n' "$attempt" "${status:-unknown}" "${progress:+ ($progress)}"

  case "$status" in
    completed|success|succeeded)
      if [[ $media_type == image ]]; then
        if ! download_image_result "$poll_body"; then
          echo "Completed image task did not contain a SourcesData URL or Base64 result" >&2
          exit 1
        fi
      else
        curl --fail-with-body --silent --show-error --location \
          --connect-timeout 15 --max-time 1800 \
          "$base_url/v1/videos/$task_id/content" \
          -H "Authorization: Bearer $SOURCESDATA_API_KEY" \
          --output "$output"
      fi
      printf 'Saved: %s\n' "$output"
      exit 0
      ;;
    failed|failure|cancelled|canceled)
      jq '{status: (.status // .data.status), error: (.error // .data.error)}' "$poll_body" >&2 || true
      exit 1
      ;;
  esac

  sleep "$(retry_after "$poll_headers")"
done

echo "Timed out waiting for task $task_id" >&2
exit 1