#!/usr/bin/env bash
# Download GGUF weights for the active stack profile into models/bonsai/
#
# Env (from .env / profile):
#   DOWNLOAD_BONSAI27=true|false
#   DOWNLOAD_QWEN=true|false
#   BONSAI_* / QWEN_* repo+file overrides
#
# Defaults: download whatever BONSAI_MAIN_FILE points at, plus optional peers.
set -euo pipefail

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
if [[ -f "$ROOT/.env" ]]; then
  set -a
  # shellcheck disable=SC1091
  source <(grep -E '^[A-Za-z_][A-Za-z0-9_]*=' "$ROOT/.env" | sed 's/\r$//') || true
  set +a
fi

OUT="${MODEL_DIR:-${BONSAI_MODELS_DIR:-$ROOT/models/bonsai}}"
# Resolve relative paths from STACK root
[[ "$OUT" != /* ]] && OUT="$ROOT/${OUT#./}"

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

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

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

hf_get() {
  local repo="$1" file="$2"
  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 repo="$1" file="$2"
  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 repo="$1" file="$2"
  if [[ -f "$file" && -s "$file" ]]; then
    echo "Already present: $file ($(du -h "$file" | cut -f1))"
    return 0
  fi
  echo "Downloading $file"
  echo "  repo: $repo"
  if have_hf; then
    hf_get "$repo" "$file"
  else
    echo "  (tip: pip install -U 'huggingface_hub[cli]')"
    curl_get "$repo" "$file"
  fi
  echo "  OK: $file ($(du -h "$file" | cut -f1))"
}

# Model definitions
BONSAI27_REPO="${BONSAI27_HF_REPO:-prism-ml/Ternary-Bonsai-27B-gguf}"
BONSAI27_FILE="${BONSAI27_MAIN_FILE:-Ternary-Bonsai-27B-Q2_0.gguf}"

QWEN_REPO="${QWEN_HF_REPO:-mradermacher/Qwen3.5-9B-abliterated-v2-MAX-GGUF}"
QWEN_FILE="${QWEN_MAIN_FILE:-Qwen3.5-9B-abliterated-v2-MAX.Q4_K_M.gguf}"

# Infer flags from active model path if not set
ACTIVE_FILE="$(basename "${BONSAI_MAIN_FILE:-${BONSAI_MODEL_PATH:-$BONSAI27_FILE}}")"
DOWNLOAD_BONSAI27="${DOWNLOAD_BONSAI27:-}"
DOWNLOAD_QWEN="${DOWNLOAD_QWEN:-}"

if [[ -z "$DOWNLOAD_BONSAI27" && -z "$DOWNLOAD_QWEN" ]]; then
  # Download the active primary file's family
  if [[ "$ACTIVE_FILE" == *"Qwen"* || "$ACTIVE_FILE" == *"qwen"* ]]; then
    DOWNLOAD_QWEN=true
    DOWNLOAD_BONSAI27=false
  else
    DOWNLOAD_BONSAI27=true
    DOWNLOAD_QWEN="${DOWNLOAD_QWEN:-false}"
  fi
fi

# Always honour DOWNLOAD_MODEL=false
if [[ "${DOWNLOAD_MODEL:-true}" == "false" || "${DOWNLOAD_MODEL:-true}" == "0" ]]; then
  echo "DOWNLOAD_MODEL=false — nothing to do"
  exit 0
fi

echo "Target directory: $OUT"
echo "Profile: ${STACK_PROFILE:-unset}  active file: $ACTIVE_FILE"
echo ""

if [[ "${DOWNLOAD_BONSAI27}" == "true" || "${DOWNLOAD_BONSAI27}" == "1" ]]; then
  download "$BONSAI27_REPO" "$BONSAI27_FILE"
  if [[ "${DOWNLOAD_VISION:-0}" == "1" || "${DOWNLOAD_VISION:-false}" == "true" ]]; then
    echo "Vision mmproj for Bonsai 27B is optional; browse HF repo if needed."
  fi
fi

if [[ "${DOWNLOAD_QWEN}" == "true" || "${DOWNLOAD_QWEN}" == "1" ]]; then
  download "$QWEN_REPO" "$QWEN_FILE"
fi

# Also download whatever BONSAI_MAIN_FILE is if different
if [[ -n "${BONSAI_HF_REPO:-}" && -n "${BONSAI_MAIN_FILE:-}" ]]; then
  if [[ ! -f "$BONSAI_MAIN_FILE" ]]; then
    download "$BONSAI_HF_REPO" "$BONSAI_MAIN_FILE"
  fi
fi

echo ""
echo "GGUF files in $OUT:"
ls -lh "$OUT"/*.gguf 2>/dev/null || ls -lh "$OUT"