#!/bin/bash
set -e

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
VENDOR_DIR="${REPO_ROOT}/vendor"

case "${DEEPINFER_ACCELERATOR:-nvidia}" in
  nvidia)
    BUILD_TARGET=CUDA
    PACKAGES=(DeepSelect FlashMLA DeepEP DeepGEMM)
    INSTALL=(deep_select flash_mla deep_ep deep_gemm)
    ;;
  ascend)
    BUILD_TARGET=ASCEND
    PACKAGES=(DeepSelect FlashMLA DeepEP-Ascend DeepGEMM-Ascend)
    INSTALL=(deep_select flash_mla deep_ep deep_gemm)
    ;;
  *)
    echo "Unsupported DEEPINFER_ACCELERATOR: '${DEEPINFER_ACCELERATOR}' (expect nvidia or ascend)" >&2
    exit 1
    ;;
esac

PLATFORM_TAG="${DEEPINFER_ACCELERATOR:-nvidia}-$(uname -m)"
WHEEL_DIR="${WHEEL_DIR:-${REPO_ROOT}/build/${PLATFORM_TAG}/vendor-wheels}"

echo "📥 Initializing submodules ..."
git -C "${REPO_ROOT}" submodule update --init

target_var() {
  case "$1" in
    DeepSelect) echo DEEP_SELECT_BUILD_TARGET_PLATFORM ;;
    FlashMLA) echo FLASH_MLA_BUILD_TARGET_PLATFORM ;;
    *) echo "" ;;
  esac
}

mkdir -p "${WHEEL_DIR}"
for pkg in "${PACKAGES[@]}"; do
  echo "🔧 Building ${pkg} (${BUILD_TARGET}) ..."
  envs=()
  var="$(target_var "${pkg}")"
  [ -n "${var}" ] && envs+=("${var}=${BUILD_TARGET}")
  [ "${pkg}" = DeepGEMM ] && envs+=(DG_FORCE_BUILD=1)
  (
    cd "${VENDOR_DIR}/${pkg}"
    git submodule update --init --recursive
    env "${envs[@]}" python setup.py bdist_wheel --dist-dir "${WHEEL_DIR}"
  )
done

echo "📦 Installing wheels from ${WHEEL_DIR} ..."
python -m pip install --no-deps --no-index --find-links "${WHEEL_DIR}" "${INSTALL[@]}"
echo "✅ vendor operators installed"