#!/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 ROCm JAX artifacts.
# Usage: ./ci/build_rocm_artifacts.sh "<artifact>"
# Supported artifact values are: jax-rocm-plugin, jax-rocm-pjrt, jax, jaxlib
#
# -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

# Clone XLA at HEAD if path to local XLA is not provided
if [[ -z "$JAXCI_XLA_GIT_DIR" && -z "$JAXCI_CLONE_MAIN_XLA" ]]; then
    export JAXCI_CLONE_MAIN_XLA=1
fi

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

# "jax" and "jaxlib" are backend neutral; they are built here so that all four
# wheels can come from one commit, stamped with one version.
allowed_artifacts=("jax-rocm-plugin" "jax-rocm-pjrt" "jax" "jaxlib")

if [[ ! " ${allowed_artifacts[*]} " =~ " ${artifact} " ]]; then
  echo "Error: Invalid artifact: $artifact. Allowed values are: ${allowed_artifacts[*]}"
  exit 1
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

override_xla_repo=""
if [[ "$JAXCI_CLONE_MAIN_XLA" == 1 ]]; then
  override_xla_repo="--bazel_options=--override_repository=xla=${JAXCI_XLA_GIT_DIR} --bazel_options=--override_module=xla=${JAXCI_XLA_GIT_DIR}"
fi

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

# Point the build ROCm at the TheRock SDK already installed in this image, so
# the wheel's ROCm SONAMEs (e.g. librocsolver) match the runtime image. Without
# this, rules_ml_toolchain silently downloads its default distro (a different,
# older ROCm) and the wheel links e.g. librocsolver.so.1 while the runtime only
# has librocsolver.so.0. rules_ml_toolchain honors ROCM_PATH ahead of its
# download path; rocm-sdk path --root is the pip-installed SDK root (no
# /opt/rocm needed). Gated on rocm-sdk so apt-ROCm images keep the default.
rocm_path_flags=()
if command -v rocm-sdk >/dev/null 2>&1; then
  rocm_path_flags=("--bazel_options=--repo_env=ROCM_PATH=$(rocm-sdk path --root)")
fi

echo "Building $artifact..."

python build/build.py build --wheels="$artifact" \
  --bazel_startup_options="--bazelrc=build/rocm/rocm.bazelrc" \
  $bazel_startup_options \
  --bazel_options=--config=rocm_release_wheel \
  --bazel_options=--config=rocm_rbe \
  --bazel_options=--remote_download_toplevel \
  --python_version=$JAXCI_HERMETIC_PYTHON_VERSION \
  --verbose --detailed_timestamped_log \
  --output_path="$JAXCI_OUTPUT_DIR" \
  $artifact_tag_flags \
  $override_xla_repo \
  "${rocm_path_flags[@]}"

# Verify manylinux compliance. "jax" is pure Python, as in ci/build_artifacts.sh.
if [[ "$artifact" != "jax" ]]; then
  ./ci/utilities/run_auditwheel.sh
fi