#
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

set -e

PYTHON_MAJOR_VERSION=${PYTHON_MAJOR_VERSION:-3}
PYTHON_MINOR_VERSION=${PYTHON_MINOR_VERSION:-8}
TARGET=${TARGET_ARCHITECTURE:-x86_64}
ROOT_PATH=${TRT_OSSPATH:-/workspace/TensorRT}
EXT_PATH=${EXT_PATH:-/tmp/external}
WHEEL_OUTPUT_DIR=${ROOT_PATH}/python/build

mkdir -p ${WHEEL_OUTPUT_DIR}
pushd ${WHEEL_OUTPUT_DIR}

# Generate tensorrt.so
cmake .. -DCMAKE_BUILD_TYPE=Release \
         -DTARGET=${TARGET} \
         -DPYTHON_MAJOR_VERSION=${PYTHON_MAJOR_VERSION} \
         -DPYTHON_MINOR_VERSION=${PYTHON_MINOR_VERSION} \
         -DEXT_PATH=${EXT_PATH} \
         -DTENSORRT_ROOT=${ROOT_PATH} \
         -DTENSORRT_MODULE=${TENSORRT_MODULE} \
         -DTENSORRT_LIBPATH=${TRT_LIBPATH} \
         -DTRT_ONNXPARSER_NAME=nvonnxparser
make -j12

# Generate wheel
TRT_MAJOR=$(awk '/^\#define TRT_MAJOR_ENTERPRISE/ {print $3}' ${ROOT_PATH}/include/NvInferVersion.h)
TRT_MINOR=$(awk '/^\#define TRT_MINOR_ENTERPRISE/ {print $3}' ${ROOT_PATH}/include/NvInferVersion.h)
TRT_PATCH=$(awk '/^\#define TRT_PATCH_ENTERPRISE/ {print $3}' ${ROOT_PATH}/include/NvInferVersion.h)
TRT_BUILD=$(awk '/^\#define TRT_BUILD_ENTERPRISE/ {print $3}' ${ROOT_PATH}/include/NvInferVersion.h)
TRT_VERSION=${TRT_MAJOR}.${TRT_MINOR}.${TRT_PATCH}.${TRT_BUILD}
TRT_MAJMINPATCH=${TRT_MAJOR}.${TRT_MINOR}.${TRT_PATCH}
PLUGIN_DISABLED_STR="0"

echo "Generating python ${PYTHON_MAJOR_VERSION}.${PYTHON_MINOR_VERSION} bindings for TensorRT ${TRT_VERSION}"

expand_vars_cp () {
    test -f ${1} || (echo "ERROR: File: ${1} does not exist!" && exit 1); \
    sed -e "s|\#\#TENSORRT_VERSION\#\#|${TRT_VERSION}|g" \
        -e "s|\#\#TENSORRT_MAJMINPATCH\#\#|${TRT_MAJMINPATCH}|g" \
        -e "s|\#\#TENSORRT_PYTHON_VERSION\#\#|${TRT_MAJMINPATCH}|g" \
        -e "s|\#\#TENSORRT_MODULE\#\#|${TENSORRT_MODULE}|g" \
        -e "s|##TENSORRT_PLUGIN_DISABLED##|${PLUGIN_DISABLED_STR}|g" \
        ${1} > ${2}
}

pushd ${ROOT_PATH}/python/packaging
for dir in $(find . -type d); do mkdir -p ${WHEEL_OUTPUT_DIR}/$dir; done
for file in $(find . -type f); do expand_vars_cp $file ${WHEEL_OUTPUT_DIR}/${file}; done
popd

# Copy required files and directories to build the bindings.
cp tensorrt/tensorrt.so bindings_wheel/tensorrt/tensorrt.so
mv tensorrt_build_backend bindings_wheel/

pushd ${WHEEL_OUTPUT_DIR}/bindings_wheel
python -m build --no-isolation --wheel --config-setting=wheel-type=binding --config-setting=python-tag=cp${PYTHON_MAJOR_VERSION}${PYTHON_MINOR_VERSION} --config-setting=plat-name=linux_${TARGET}
popd