#!/usr/bin/env bash
# Download BGE-Reranker-v2-m3 FP16 GGUF and apply YAML profile.
# Usage: ./scripts/download-reranker.sh
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
STACK_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"

cd "${STACK_DIR}"

if [[ -f .env ]]; then
  set -a
  # shellcheck disable=SC1091
  source .env
  set +a
fi

DATA_ROOT="${DATA_ROOT:-/data}"
MODELS_DIR="${DATA_ROOT}/apps/localai/models"
GGUF_NAME="bge-reranker-v2-m3-FP16.gguf"
GGUF_URL="https://huggingface.co/gpustack/bge-reranker-v2-m3-GGUF/resolve/main/${GGUF_NAME}"
YAML_SRC="${STACK_DIR}/profiles/bge-reranker-v2-m3-FP16-rerank.yaml.example"
YAML_DST="${MODELS_DIR}/${GGUF_NAME}.yaml"

"${SCRIPT_DIR}/ensure-dirs.sh" "${DATA_ROOT}"

echo "=== BGE-Reranker-v2-m3 download ==="
echo "Target: ${MODELS_DIR}/${GGUF_NAME}"
echo ""

if [[ -f "${MODELS_DIR}/${GGUF_NAME}" ]]; then
  echo "GGUF already exists — skipping download"
else
  if command -v wget &>/dev/null; then
    wget -c -O "${MODELS_DIR}/${GGUF_NAME}.partial" "${GGUF_URL}"
    mv "${MODELS_DIR}/${GGUF_NAME}.partial" "${MODELS_DIR}/${GGUF_NAME}"
  else
    curl -fL -C - -o "${MODELS_DIR}/${GGUF_NAME}.partial" "${GGUF_URL}"
    mv "${MODELS_DIR}/${GGUF_NAME}.partial" "${MODELS_DIR}/${GGUF_NAME}"
  fi
fi

echo ""
echo "=== Applying YAML profile ==="
if cp "${YAML_SRC}" "${YAML_DST}" 2>/dev/null; then
  chmod 644 "${YAML_DST}" "${MODELS_DIR}/${GGUF_NAME}" 2>/dev/null || true
elif docker ps --format '{{.Names}}' | grep -qx localai; then
  echo "Host copy failed (permissions) — writing via docker exec localai"
  docker exec -i localai sh -c "cat > /models/${GGUF_NAME}.yaml" < "${YAML_SRC}"
else
  echo "ERROR: cannot write ${YAML_DST} (permission denied) and localai container not running"
  exit 1
fi

ls -lh "${MODELS_DIR}/${GGUF_NAME}" "${YAML_DST}"

echo ""
echo "=== Done ==="
echo "Restart LocalAI to load the model:"
echo "  cd ${STACK_DIR} && docker compose --profile localai restart localai"
echo ""
echo "Smoke test:"
echo '  curl -s http://127.0.0.1:${LOCALAI_PORT:-8070}/v1/rerank \'
echo '    -H "Authorization: Bearer <LOCALAI_API_KEY>" \'
echo '    -H "Content-Type: application/json" \'
echo '    -d '"'"'{"model":"bge-reranker-v2-m3-FP16.gguf","query":"panda","documents":["hi","it is a bear","The giant panda is a bear species endemic to China."],"top_n":2}'"'"