# Setup hook for ascend-toolkit
# This is sourced when ascend-toolkit is in buildInputs of another derivation
addAscendEnv() {
    local toolkit="${1:?ascend-toolkit path required}"

    # Core environment variables
    export ASCEND_TOOLKIT_HOME="${toolkit}"
    export ASCEND_HOME_PATH="${toolkit}"
    export ASCEND_OPP_PATH="${toolkit}/opp"
    export ASCEND_AICPU_PATH="${toolkit}"
    export TOOLCHAIN_HOME="${toolkit}/toolkit"

    # PATH
    addToSearchPath PATH "${toolkit}/bin"
    addToSearchPath PATH "${toolkit}/tools/ccec_compiler/bin"
    addToSearchPath PATH "${toolkit}/tools/profiler/bin"
    addToSearchPath PATH "${toolkit}/tools/ascend_system_advisor/asys"
    addToSearchPath PATH "${toolkit}/tools/show_kernel_debug_data"
    addToSearchPath PATH "${toolkit}/tools/msobjdump"

    # LD_LIBRARY_PATH
    addToSearchPath LD_LIBRARY_PATH "${toolkit}/lib64"
    addToSearchPath LD_LIBRARY_PATH "${toolkit}/lib64/plugin/opskernel"
    addToSearchPath LD_LIBRARY_PATH "${toolkit}/lib64/plugin/nnengine"
    addToSearchPath LD_LIBRARY_PATH "${toolkit}/opp/built-in/op_impl/ai_core/tbe/op_tiling/lib/linux/x86_64"
    addToSearchPath LD_LIBRARY_PATH "${toolkit}/tools/aml/lib64"
    addToSearchPath LD_LIBRARY_PATH "${toolkit}/tools/aml/lib64/plugin"
    addToSearchPath LD_LIBRARY_PATH "${toolkit}/devlib"

    # PYTHONPATH
    addToSearchPath PYTHONPATH "${toolkit}/python/site-packages"
    addToSearchPath PYTHONPATH "${toolkit}/opp/built-in/op_impl/ai_core/tbe"

    # CMAKE_PREFIX_PATH
    addToSearchPath CMAKE_PREFIX_PATH "${toolkit}/lib64/cmake"
    addToSearchPath CMAKE_PREFIX_PATH "${toolkit}/toolkit/tools/tikicpulib/lib/cmake"
}

if [ -n "${ascend-toolkit:-}" ]; then
    addAscendEnv "${ascend-toolkit}"
fi