# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

# limitations under the License.
# ==============================================================================
source ci/official/envs/linux_arm64
TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config  cross_compile_linux_arm64"

# limitations under the License.
# ==============================================================================
source ci/official/envs/linux_x86_build
TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config release_cpu_linux --config=tpu"
TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu
TFCI_BUILD_PIP_PACKAGE_ARGS="--repo_env=WHEEL_NAME=tensorflow_tpu"
TFCI_LIB_SUFFIX="-tpu-linux-x86_64"

  echo "Error: Generated wheel is too big! Limit is $TFCI_WHL_SIZE_LIMIT"
  echo '(search for TFCI_WHL_SIZE_LIMIT to change it)'
  ls -sh *.whl
  exit 2
fi

# Quick install checks
venv=$(mktemp -d)
"python${TFCI_PYTHON_VERSION}" -m venv "$venv"
python="$venv/bin/python3"
# TODO(b/366266944) Remove the check after tf docker image upgrade for NumPy 2
# and numpy 1 support is dropped b/361369076.
if [[ "$TFCI_WHL_NUMPY_VERSION" == 1 ]]; then

# Temporary envs file while we migrate to the new docker image.
# Once the migration is complete, this file will be removed.
source ci/official/envs/linux_x86_build
TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config release_gpu_linux"
TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda
TFCI_BUILD_PIP_PACKAGE_ARGS="--repo_env=WHEEL_NAME=tensorflow"
TFCI_DOCKER_ARGS="--gpus all"
TFCI_LIB_SUFFIX="-gpu-linux-x86_64"

TFCI_MACOS_TWINE_INSTALL_ENABLE=
TFCI_MACOS_UPGRADE_PYENV_ENABLE=
TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=
TFCI_NVIDIA_SMI_ENABLE=
TFCI_OUTPUT_DIR=
TFCI_PYCPP_SWAP_TO_BUILD_ENABLE=
TFCI_PYTHON_VERIFY_PIP_INSTALL_ARGS=
TFCI_PYTHON_VERSION=
TFCI_WHL_AUDIT_ENABLE=
TFCI_WHL_AUDIT_PLAT=
TFCI_WHL_BAZEL_TEST_ENABLE=
TFCI_WHL_IMPORT_TEST_ENABLE=1
TFCI_WHL_NUMPY_VERSION=2
TFCI_WHL_SIZE_LIMIT_ENABLE=

fi

if [[ "$TFCI_WHL_BAZEL_TEST_ENABLE" == 1 ]]; then
  tfrun bazel test $TFCI_BAZEL_COMMON_ARGS $TFCI_BUILD_PIP_PACKAGE_ARGS --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config="${TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX}_wheel_test"