#!/usr/bin/env bash
# .devcontainer/postCreateCommand.sh
# Faiss NPU Dev Container 初始化脚本（基于官方 CANN 预制镜像）
# 在容器首次创建时执行一次：安装编译依赖、升级 CMake、验证 NPU 环境。

set -e

echo "=========================================="
echo "Faiss NPU Dev Container post-create setup..."
echo "=========================================="

# 1. 切换 apt 源为华为云镜像（若尚未切换，国内网络加速）
if ! grep -q "mirrors.huaweicloud.com" /etc/apt/sources.list 2>/dev/null; then
    cp -a /etc/apt/sources.list /etc/apt/sources.list.bak
    sed -i "s@http://.*archive.ubuntu.com@https://mirrors.huaweicloud.com@g" /etc/apt/sources.list
    sed -i "s@http://.*security.ubuntu.com@https://mirrors.huaweicloud.com@g" /etc/apt/sources.list
fi

# 2. 安装系统依赖（幂等：已安装则跳过）
if ! dpkg -s libopenblas-dev >/dev/null 2>&1; then
    echo "Installing system dependencies..."
    apt-get update
    apt-get install -y --no-install-recommends \
        libopenblas-dev libopenblas0 swig python3-dev python3-numpy gawk \
        libgtest-dev libomp-dev
    rm -rf /var/lib/apt/lists/*
else
    echo "System dependencies already installed."
fi

# 3. 升级 CMake 至 >= 3.24（CANN 镜像自带 3.22 不满足要求）
CMAKE_VER=$(cmake --version 2>/dev/null | head -n1 | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' || true)
if [ -z "${CMAKE_VER}" ] || dpkg --compare-versions "${CMAKE_VER}" lt "3.24.0"; then
    echo "Upgrading CMake to 3.26.4..."
    ARCH=$(uname -m)
    case "${ARCH}" in
        x86_64) CMAKE_ARCH="x86_64" ;;
        aarch64) CMAKE_ARCH="aarch64" ;;
        *) echo "Unsupported arch: ${ARCH}" >&2; exit 1 ;;
    esac
    wget -q "https://github.com/Kitware/CMake/releases/download/v3.26.4/cmake-3.26.4-linux-${CMAKE_ARCH}.tar.gz" \
        -O /tmp/cmake.tar.gz
    tar -xzf /tmp/cmake.tar.gz -C /usr/local
    ln -sf "/usr/local/cmake-3.26.4-linux-${CMAKE_ARCH}/bin/cmake" /usr/local/bin/cmake
    ln -sf "/usr/local/cmake-3.26.4-linux-${CMAKE_ARCH}/bin/ctest" /usr/local/bin/ctest
    rm -f /tmp/cmake.tar.gz
    hash -r
fi
echo "CMake   : $(cmake --version | head -n1)"

# 4. Python 依赖（固定 numpy 版本 + 构建/测试所需包，华为云 pip 镜像）
pip3 install --no-cache-dir --upgrade \
    "numpy==1.26.4" packaging scipy pytest \
    -i https://repo.huaweicloud.com/repository/pypi/simple

echo "Python  : $(python3 --version)"
echo "GCC     : $(gcc --version | head -n1)"
echo "CANN    : ${ASCEND_HOME_PATH:-/usr/local/Ascend/ascend-toolkit/latest}"

# 5. 验证 NPU 环境
if command -v npu-smi &> /dev/null; then
    echo "NPU environment detected:"
    npu-smi info || true
else
    echo "Warning: npu-smi not found. NPU may not be available."
fi

echo ""
echo "--- 构建步骤 ---"
echo "1) 编译 NPU 算子:"
echo "   cd ${PWD}/faiss/npu/ops && bash ops_build.sh"
echo "   # 若芯片为 A3，使用: COMPUTE_UNIT=ascend910_93 bash ops_build.sh"
echo ""
echo "2) 编译 faiss_npu:"
echo "   mkdir -p build && cd build"
echo "   cmake .. -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF \\"
echo "            -DFAISS_ENABLE_GPU=OFF -DFAISS_ENABLE_NPU=ON -DBUILD_TESTING=OFF"
echo "   make -j\$(nproc)"
echo ""
echo "3) 打包安装:"
echo "   cd build/faiss/python && python3 setup.py bdist_wheel"
echo "   pip3 install dist/faiss_ascend-*.whl"
echo "=========================================="