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

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

# shellcheck disable=SC1091
source "${SCRIPT_DIR}/catalog-lib.sh"

cd "${STACK_DIR}"

if [[ $# -lt 1 ]]; then
  echo "Usage: $0 <catalog-model-id>"
  echo ""
  echo "Available models:"
  catalog_list_ids | sed 's/^/  /'
  exit 1
fi

MODEL_ID="$1"

if ! catalog_list_ids | grep -qx "${MODEL_ID}"; then
  echo "ERROR: Unknown model id: ${MODEL_ID}"
  echo "Run: ./scripts/list-models.sh"
  exit 1
fi

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

DATA_ROOT="${DATA_ROOT:-/data}"
catalog_ensure_dirs "${DATA_ROOT}"

RUNTIME="$(catalog_get_field "${MODEL_ID}" runtime)"
NAME="$(catalog_get_field "${MODEL_ID}" name)"

echo "=== Download: ${NAME} (${MODEL_ID}) ==="
echo "Runtime: ${RUNTIME}"
echo ""

case "${RUNTIME}" in
  vllm)
    HF_MODEL="$(catalog_get_field "${MODEL_ID}" hf_model)"
    echo "Target: Hugging Face cache at ${DATA_ROOT}/apps/vllm/huggingface"
    echo "Model:  ${HF_MODEL}"
    echo ""

    HF_ENV=()
    if [[ -n "${HF_TOKEN:-}" ]]; then
      HF_ENV=(-e "HF_TOKEN=${HF_TOKEN}")
    fi

    docker run --rm \
      "${HF_ENV[@]}" \
      -v "${DATA_ROOT}/apps/vllm/huggingface:/root/.cache/huggingface" \
      ghcr.io/huggingface/huggingface-cli:latest \
      download "${HF_MODEL}"
  ;;
  llamacpp)
    DOWNLOAD_URL="$(catalog_get_field "${MODEL_ID}" download_url)"
    LOCAL_PATH="$(catalog_get_field "${MODEL_ID}" local_path)"
    LOCAL_DIR="$(catalog_get_field "${MODEL_ID}" local_dir)"
    GGUF_FILE="$(catalog_get_field "${MODEL_ID}" gguf_file)"
    MMPROJ_URL="$(catalog_get_field "${MODEL_ID}" mmproj_url)"
    MMPROJ_FILE="$(catalog_get_field "${MODEL_ID}" mmproj_file)"

    if [[ -n "${LOCAL_PATH}" ]]; then
      DEST="${LOCAL_PATH}"
      mkdir -p "$(dirname "${DEST}")"
      if [[ -f "${DEST}" ]]; then
        echo "Already exists: ${DEST}"
      else
        echo "Downloading ${DEST} ..."
        wget -c -O "${DEST}" "${DOWNLOAD_URL}"
      fi
    elif [[ -n "${LOCAL_DIR}" && -n "${GGUF_FILE}" ]]; then
      mkdir -p "${LOCAL_DIR}"
      DEST="${LOCAL_DIR}/${GGUF_FILE}"
      if [[ -f "${DEST}" ]]; then
        echo "Already exists: ${DEST}"
      else
        echo "Downloading ${DEST} ..."
        wget -c -O "${DEST}" "${DOWNLOAD_URL}"
      fi
      if [[ -n "${MMPROJ_URL}" && -n "${MMPROJ_FILE}" ]]; then
        MMPROJ_DEST="${LOCAL_DIR}/${MMPROJ_FILE}"
        if [[ -f "${MMPROJ_DEST}" ]]; then
          echo "Already exists: ${MMPROJ_DEST}"
        else
          echo "Downloading ${MMPROJ_DEST} ..."
          wget -c -O "${MMPROJ_DEST}" "${MMPROJ_URL}"
        fi
      fi
    else
      echo "ERROR: Catalog entry missing local_path or local_dir/gguf_file"
      exit 1
    fi

    echo ""
    echo "GGUF saved for future llama.cpp stack (stacks/llamacpp/README.md)."
    echo "Standard vLLM cannot load .gguf files."
  ;;
  *)
    echo "ERROR: Unknown runtime: ${RUNTIME}"
    exit 1
    ;;
esac

echo ""
echo "Done."