#!/usr/bin/env bash
# Download Ternary-Bonsai-27B GGUF (and optional vision / drafter packs)
# into ./models/bonsai for the docker-compose bonsai service.
#
# Requires: ~8–10 GB free disk for the base model.
# Prefer:   pip install -U "huggingface_hub[cli]"   # provides `hf`
# Optional: HF_TOKEN / BONSAI_TOKEN if the repo ever requires auth.
set -euo pipefail

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
OUT="${MODEL_DIR:-$ROOT/models/bonsai}"
REPO="${BONSAI_HF_REPO:-prism-ml/Ternary-Bonsai-27B-gguf}"
# Group-128 Q2_0 — needs PrismML llama.cpp (what our Docker image ships)
MAIN_FILE="${BONSAI_MAIN_FILE:-Ternary-Bonsai-27B-Q2_0.gguf}"
DOWNLOAD_VISION="${DOWNLOAD_VISION:-1}"
DOWNLOAD_DRAFTER="${DOWNLOAD_DRAFTER:-0}"

mkdir -p "$OUT"
cd "$OUT"

export HF_TOKEN="${HF_TOKEN:-${BONSAI_TOKEN:-}}"

have_hf() {
  command -v hf >/dev/null 2>&1 || command -v huggingface-cli >/dev/null 2>&1
}

hf_get() {
  local file="$1"
  if command -v hf >/dev/null 2>&1; then
    hf download "$REPO" "$file" --local-dir .
  else
    huggingface-cli download "$REPO" "$file" --local-dir .
  fi
}

curl_get() {
  local file="$1"
  local url="https://huggingface.co/${REPO}/resolve/main/${file}"
  local auth=()
  if [ -n "${HF_TOKEN:-}" ]; then
    auth=(-H "Authorization: Bearer ${HF_TOKEN}")
  fi
  curl -fL --retry 5 --continue-at - "${auth[@]}" "$url" -o "${file}.partial"
  mv "${file}.partial" "$file"
}

download() {
  local file="$1"
  if [ -f "$file" ] && [ -s "$file" ]; then
    echo "Already present: $file ($(du -h "$file" | cut -f1))"
    return 0
  fi
  echo "Downloading $file ..."
  if have_hf; then
    hf_get "$file"
  else
    echo "  (tip: install huggingface_hub for more reliable downloads: pip install -U 'huggingface_hub[cli]')"
    curl_get "$file"
  fi
  echo "  OK: $file ($(du -h "$file" | cut -f1))"
}

echo "Target directory: $OUT"
echo "Repo:             $REPO"
echo ""

download "$MAIN_FILE"

if [ "$DOWNLOAD_VISION" = "1" ]; then
  echo ""
  echo "Optional vision projector (mmproj) — attempting known filenames..."
  _got=0
  for f in \
    Ternary-Bonsai-27B-mmproj-HQQ4.gguf \
    Ternary-Bonsai-27B-mmproj-f16.gguf \
    mmproj-Ternary-Bonsai-27B.gguf \
    ; do
    if [ -f "$f" ]; then
      echo "Already present: $f"
      _got=1
      break
    fi
    if have_hf; then
      if hf_get "$f" 2>/dev/null; then
        _got=1
        break
      fi
    fi
  done
  if [ "$_got" != "1" ]; then
    echo "Vision mmproj not found automatically (optional for text-only)."
    echo "  List files: https://huggingface.co/${REPO}/tree/main"
    echo "  Place any *mmproj*.gguf into $OUT and restart bonsai."
  fi
fi

if [ "$DOWNLOAD_DRAFTER" = "1" ]; then
  echo ""
  echo "Optional DSpark drafter (speculative decoding; mainly helps CUDA)..."
  for f in Ternary-Bonsai-27B-dspark-Q4_1.gguf; do
    download "$f" || echo "  skipped: $f"
  done
fi

echo ""
echo "Done. GGUF files in $OUT:"
ls -lh "$OUT"/*.gguf 2>/dev/null || ls -lh "$OUT"
echo ""
echo "Next:"
echo "  docker compose up -d --build bonsai open-webui searxng"
echo "  Open WebUI → http://localhost:3000"
echo "  Raw API   → http://localhost:8081/v1/models"