#!/bin/bash
# Copyright 2026 The JAX Authors.
##
# 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.
# ==============================================================================
# Build JAX artifacts.
# Usage: ./ci/build_oneapi_artifacts.sh "<artifact>"
# Supported artifact values are: jax-oneapi-plugin, jax-oneapi-pjrt
# E.g: ./ci/build_oneapi_artifacts.sh "jax-oneapi-plugin" or ./ci/build_oneapi_artifacts.sh "jax-oneapi-pjrt"
#
# -e: abort script if one command fails
# -u: error if undefined variable used
# -x: log all commands
# -o history: record shell history
# -o allexport: export all functions and variables to be available to subscripts
set -exu -o history -o allexport

artifact="$1"

# Source default JAXCI environment variables.
source ci/envs/default.env

# Set up the build environment.
source "ci/utilities/setup_build_environment.sh"

allowed_artifacts=("jax-oneapi-plugin" "jax-oneapi-pjrt")

if [[ ! " ${allowed_artifacts[*]} " =~ " ${artifact} " ]]; then
  echo "Error: Invalid artifact: $artifact. Allowed values are: ${allowed_artifacts[*]}"
  exit 1
fi

os=$(uname -s | awk '{print tolower($0)}')
arch=$(uname -m)

bazel_startup_options=""
if [[ -n "${JAXCI_BAZEL_OUTPUT_BASE}" ]]; then
  bazel_startup_options="--output_base=${JAXCI_BAZEL_OUTPUT_BASE}"
fi

# Determine the Bazel flags that stamp the wheel's version.
source "ci/utilities/set_artifact_tag_flags.sh"

if [[ "$JAXCI_HERMETIC_PYTHON_VERSION" == *t && "$JAXCI_HERMETIC_PYTHON_VERSION" != *"-ft" ]]; then
  JAXCI_HERMETIC_PYTHON_VERSION=${JAXCI_HERMETIC_PYTHON_VERSION%t}-ft
fi

# Figure out the bazelrc config to use. We will use one of the "rbe_"/"ci_"
# flags in the .bazelrc depending upon the platform we are building for.
bazelrc_config="${os}_${arch}"

# On platforms with no RBE support, we can use the Bazel remote cache. Set
# it to be empty by default to avoid unbound variable errors.
bazel_remote_cache=""

if [[ "$JAXCI_BUILD_ARTIFACT_WITH_RBE" == 1 ]]; then
  bazelrc_config="rbe_${bazelrc_config}"
  bazel_cpu_pool_config="--bazel_options=--config=rbe_cpu_pool"
else
  bazelrc_config="ci_${bazelrc_config}"
  bazel_cpu_pool_config=""

  # Set remote cache flags. Pushes to the cache bucket is limited to JAX's
  # CI system.
  if [[ "$JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE" == 1 ]]; then
    bazel_remote_cache="--bazel_options=--config=public_cache_push"
  else
    bazel_remote_cache="--bazel_options=--config=public_cache"
  fi
fi

oneapi_version_flag="--oneapi_version=$JAXCI_ONEAPI_VERSION"

# Build the artifact.
python build/build.py build --wheels="$artifact" \
  --bazel_options=--config="$bazelrc_config" $bazel_remote_cache \
  $bazel_cpu_pool_config \
  --bazel_startup_options="$bazel_startup_options" \
  --python_version=$JAXCI_HERMETIC_PYTHON_VERSION \
  $oneapi_version_flag \
  --verbose --detailed_timestamped_log \
  --output_path="$JAXCI_OUTPUT_DIR" \
  $artifact_tag_flags

# If building release artifacts, we also build a release candidate ("rc")
# tagged wheel. A build using JAXCI_WHEEL_VERSION_SUFFIX already has its rc
# wheel from above.
if [[ "$JAXCI_ARTIFACT_TYPE" == "release" && -n "$JAXCI_WHEEL_RC_VERSION" ]]; then
  python build/build.py build --wheels="$artifact" \
    --bazel_options=--config="$bazelrc_config" $bazel_remote_cache \
    $bazel_cpu_pool_config \
    --bazel_startup_options="$bazel_startup_options" \
    --python_version=$JAXCI_HERMETIC_PYTHON_VERSION \
    $oneapi_version_flag \
    --verbose --detailed_timestamped_log \
    --output_path="$JAXCI_OUTPUT_DIR" \
    $artifact_tag_flags --bazel_options=--repo_env=ML_WHEEL_VERSION_SUFFIX="$JAXCI_WHEEL_RC_VERSION"
fi

./ci/utilities/run_auditwheel.sh