#!/usr/bin/env bash
# ===========================================================================
# SearchAI Inference Server — on-demand model fetcher.
# Downloads GGUF models (and TTS voices) into the models directory. Same model
# files as the Java/Rust engines; default source is the shared SearchBlox S3
# bucket, with an anonymous public-URL fallback when the aws CLI is absent.
#
# Groups:
#   4b     q35-4b.gguf + q35-4b-mmproj.gguf         ~3.4 GB   text + VISION (default test model)
#   asr    qwen3-asr-1.7b.gguf + mmproj-qwen3-asr    ~2.5 GB   speech -> text
#   tts    qwen-talker-1.7b.gguf + qwen-tts-codec    ~2.7 GB   text -> speech
#   voice  tts-voices/*.wav                           small    TTS preset voices (if published)
#   2b     q35-2b.gguf + q35-2b-mmproj.gguf          ~1.9 GB   smaller text + vision
#   9b     q35-9b.gguf + q35-9b-mmproj.gguf          ~6.6 GB   best text + vision
#   image  qwen-image-edit-2511 (Q4) + text encoder + vae  ~30 GB  instruction image editing (GPU recommended)
#   all    4b + asr + tts + voice (image NOT included — fetch explicitly)
#
# Usage:
#   sudo ./fetch-models.sh --dest /var/lib/searchai/models 4b asr tts voice
#   ./fetch-models.sh --list
# ===========================================================================
set -euo pipefail
BUCKET="${SEARCHAI_MODELS_BUCKET:-s3://searchai-inference-server/models}"
DEST="${MODELS_DIR:-$(cd "$(dirname "$0")" && pwd)}"

files_for() { case "$1" in
  4b)    echo "q35-4b.gguf q35-4b-mmproj.gguf" ;;
  2b)    echo "q35-2b.gguf q35-2b-mmproj.gguf" ;;
  0.8b)  echo "q35-0.8b.gguf" ;;  # text-only (no mmproj published)
  9b)    echo "q35-9b.gguf q35-9b-mmproj.gguf" ;;
  asr)   echo "qwen3-asr-1.7b.gguf mmproj-qwen3-asr.gguf" ;;
  tts)   echo "qwen-talker-1.7b.gguf qwen-tts-codec.gguf" ;;
  image) echo "qwen-image-edit-2511-Q4_K_M.gguf qwen_2.5_vl_7b.safetensors qwen_image_vae.safetensors" ;;
  voice) echo "" ;;  # handled specially (directory of WAVs)
  all)   echo "$(files_for 4b) $(files_for asr) $(files_for tts)" ;;
  *)     echo "" ;;
esac; }

ARGS=(); while [ $# -gt 0 ]; do case "$1" in
  --dest) DEST="$2"; shift 2 ;;
  --list) sed -n '/^# Groups:/,/^# Usage:/p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
  *) ARGS+=("$1"); shift ;;
esac; done
[ ${#ARGS[@]} -eq 0 ] && { echo "usage: $0 [--dest DIR] <4b|asr|tts|voice|2b|9b|image|all> ...  (--list for details)"; exit 1; }

have_aws=0; command -v aws >/dev/null && have_aws=1
# public HTTPS fallback (CDN-cached); aws-CLI path stays S3-direct (free in-region)
PUBURL="${SEARCHAI_MODELS_URL:-https://inference-server.searchblox.com/models}"
mkdir -p "$DEST"
get() { # get <relpath>
  local f="$1" out="$DEST/$1"
  [ -f "$out" ] && { echo "  have $f"; return; }
  mkdir -p "$(dirname "$out")"
  echo "  fetching $f -> $out"
  # aws (free in-region on EC2) when it works; else public HTTPS
  if [ "$have_aws" = 1 ] && aws s3 cp "$BUCKET/$f" "$out" --only-show-errors 2>/dev/null; then :
  else
    if [ -t 1 ]; then curl -fL --progress-bar "$PUBURL/$f" -o "$out"
    else curl -fsSL "$PUBURL/$f" -o "$out"; fi
  fi
}

for g in "${ARGS[@]}"; do
  if [ "$g" = voice ]; then
    echo "voice: TTS preset WAVs -> $DEST/tts-voices/"
    if [ "$have_aws" = 1 ] && aws s3 ls "$BUCKET/tts-voices/" >/dev/null 2>&1 \
       && [ -n "$(aws s3 ls "$BUCKET/tts-voices/" 2>/dev/null)" ]; then
      aws s3 cp "$BUCKET/tts-voices/" "$DEST/tts-voices/" --recursive --only-show-errors
    else
      echo "  WARNING: no TTS preset voices published at $BUCKET/tts-voices/ —"
      echo "           TTS still works by passing reference audio in the request; or"
      echo "           drop your own <name>.wav files into $DEST/tts-voices/."
    fi
    continue
  fi
  fs="$(files_for "$g")"
  [ -z "$fs" ] && { echo "unknown group: $g (try --list)"; exit 1; }
  echo "$g:"; for f in $fs; do get "$f"; done
done
echo "done -> $DEST"