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

REPO_ROOT=${REPO_ROOT:-/data/vllm-ascend-v0.18.0}
CANN_ENV=${CANN_ENV:-/usr/local/Ascend/cann-8.5.1/cann-8.5.1/set_env.sh}
VENV_ACTIVATE=${VENV_ACTIVATE:-/data/vllm-ascend-v0.18.0/.venv/bin/activate}
INSTALLER_DIR=${INSTALLER_DIR:-/data/installers/cann-8.5.1}
TORCH_NPU_INDEX=${TORCH_NPU_INDEX:-https://mirrors.huaweicloud.com/repository/pypi/simple}
TRITON_WHEEL_NAME=triton_ascend-3.2.0.dev20260322-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
TRITON_WHEEL_URL=${TRITON_WHEEL_URL:-https://vllm-ascend.obs.cn-north-4.myhuaweicloud.com/vllm-ascend/${TRITON_WHEEL_NAME}}
TRITON_WHEEL=${INSTALLER_DIR}/${TRITON_WHEEL_NAME}

if [[ "$(realpath "${REPO_ROOT}")" != /data/vllm-ascend-v0.18.0 ]]; then
    echo "REPO_ROOT must resolve to /data/vllm-ascend-v0.18.0" >&2
    exit 2
fi
[[ -f "${CANN_ENV}" ]]
[[ -f "${VENV_ACTIVATE}" ]]

source "${CANN_ENV}"
source "${VENV_ACTIVATE}"
mkdir -p "${INSTALLER_DIR}"

python -m pip install \
    --no-deps \
    --force-reinstall \
    --index-url "${TORCH_NPU_INDEX}" \
    torch-npu==2.9.0

curl --fail --location --continue-at - \
    --output "${TRITON_WHEEL}" "${TRITON_WHEEL_URL}"
python -m pip uninstall -y triton triton-ascend || true
python -m pip install --no-deps "${TRITON_WHEEL}"

# Rebuild all CANN-bound artifacts after switching toolkit versions.
rm -rf \
    "${REPO_ROOT}/csrc/build" \
    "${REPO_ROOT}/build" \
    "${REPO_ROOT}/vllm_ascend/_cann_ops_custom"
rm -f \
    "${REPO_ROOT}"/vllm_ascend/vllm_ascend_C*.so \
    "${REPO_ROOT}"/vllm_ascend/libvllm_ascend_kernels.so

export VLLM_TARGET_DEVICE=ascend
export CMAKE_ARGS="-DPython_EXECUTABLE=${VIRTUAL_ENV}/bin/python -DPython_ROOT_DIR=${VIRTUAL_ENV}"
python -m pip install --no-build-isolation --no-deps -e "${REPO_ROOT}"

python - <<'PY'
import importlib.metadata

triton_version = importlib.metadata.version("triton-ascend")
torch_npu_version = importlib.metadata.version("torch-npu")
if triton_version != "3.2.0.dev20260322":
    raise RuntimeError(
        f"expected triton-ascend 3.2.0.dev20260322, got {triton_version}"
    )
if torch_npu_version != "2.9.0":
    raise RuntimeError(f"expected torch-npu 2.9.0, got {torch_npu_version}")
print(f"triton-ascend={triton_version}")
print(f"torch-npu={torch_npu_version}")
PY
