# All commands run relative to this directory
cd "$(dirname "${BASH_SOURCE[0]}")"

mv BUILD.bazel BUILD

SUPPORTED_VERSIONS=("3_9" "3_10" "3_11" "3_12")

for VERSION in "${SUPPORTED_VERSIONS[@]}"
do
  touch "requirements_lock_$VERSION.txt"
  bazel run \
    --experimental_convenience_symlinks=ignore \
    --enable_bzlmod=false \
    //:requirements_"$VERSION".update
  sed -i '/^#/d' requirements_lock_"$VERSION".txt

  # Install the necessary Python, unless it's already present
  pyenv install -s "$TFCI_PYTHON_VERSION"
  pyenv local "$TFCI_PYTHON_VERSION"
  # Do a sanity check to make sure that we using the correct Python version
  python --version
fi

# TFCI Mac VM images do not have twine installed by default so we need to
# install it manually. We use Twine in nightly builds to upload Python packages
# to PyPI.

# ==============================================================================
#
# Re-direct all links in $1 that are relative to be canonical

BASE="$1"
find "${BASE}" -type l | \
  while read l ; do
    if [[ "$(readlink "$l")" == \.\./* ]]; then
      CANONICAL="$(readlink "$l")";
      rm "$l";
      ln -s "${CANONICAL}" "$l"
    fi

# if there is a change in requirements.in then all lock files will be updated
# accordingly.

mv BUILD.bazel BUILD

SUPPORTED_VERSIONS=("3_9" "3_10" "3_11" "3_12")

for VERSION in "${SUPPORTED_VERSIONS[@]}"
do
  cp ../../../requirements_lock_"$VERSION".txt "requirements_lock_"$VERSION".txt"
  bazel run \
    --experimental_convenience_symlinks=ignore \
    --enable_bzlmod=false \
    //:requirements_"$VERSION"_release.update

# Check and rename wheels with auditwheel. Inserts the platform tags like
# "manylinux_xyz" into the wheel filename.
set -euxo pipefail

for wheel in /tf/pkg/*.whl; do
  echo "Checking and renaming $wheel..."
  time python3 -m auditwheel repair --plat manylinux2014_aarch64 "$wheel" --wheel-dir /tf/pkg 2>&1 | tee check.txt

  # We don't need the original wheel if it was renamed

$VERSION-distutils
EOF
/setup.packages.sh pythons.txt

# Re-link pyconfig.h from aarch64-linux-gnu into the devtoolset directory
# for any Python version present
pushd /usr/include/aarch64-linux-gnu
for f in $(ls | grep python); do
  # set up symlink for devtoolset-10
  rm -f /dt10/usr/include/aarch64-linux-gnu/$f
  ln -s /usr/include/aarch64-linux-gnu/$f /dt10/usr/include/aarch64-linux-gnu/$f
done
popd

# Python 3.10 include headers fix:

# link in the whole aarch64-linux-gnu folder, as otherwise we're overlaying
# system gcc paths that we do not want to find.
# TODO(klimek): Automate linking in all non-gcc / non-kernel include
# directories.
mkdir -p "${TARGET}/usr/include/aarch64-linux-gnu"
PYTHON_VERSIONS=("python3.8" "python3.9" "python3.10" "python3.11")
for v in "${PYTHON_VERSIONS[@]}"; do
  ln -s "/usr/local/include/${v}" "${TARGET}/usr/include/aarch64-linux-gnu/${v}"

  else
    TAG="pr-${KOKORO_GITHUB_PULL_REQUEST_NUMBER}"
  fi
fi

# Build for both JAX and TF usage.  We do these in one place because they share
# almost all of the same cache layers
export DOCKER_BUILDKIT=1
for target in jax tf; do
  IMAGE="gcr.io/tensorflow-sigs/build-arm64:$target-$TAG"
  docker pull "$IMAGE" || true
  # Due to some flakiness of resources pulled in the build, allow the docker

# limitations under the License.
# ==============================================================================
#
# Given an RPM spec file $1, apply its patches.

SPEC="$1"
grep '%patch' "${SPEC}" |while read cmd ; do
  N=$(echo "${cmd}" |sed 's,%patch\([0-9]\+\).*,\1,')
  file=$(grep "Patch$N:" "${SPEC}" |sed 's,.*: ,,')
  parg=$(echo "${cmd}" |sed 's,.*\(-p[0-9]\).*,\1,')
  if [[ ! "${file}" =~ doxygen && "${cmd}" != \#* ]]; then

#
# 2. RUN ANY OTHER SCRIPT AND ENV WITH NO SIDE EFFECTS (NO UPLOADS)
#    To use:
#       export TFCI=ci/official/envs/env_goes_here
#       export TF_ANY_SCRIPT=ci/official/wheel.sh
#       ./any.sh
#
# 3. DO THE SAME WITH A LOCAL CACHE OR RBE:
#       export TF_ANY_EXTRA_ENV=ci/official/envs/public_cache,ci/official/envs/disk_cache
#       ...
#       ./any.sh
#     or
#       export TF_ANY_EXTRA_ENV=ci/official/envs/local_rbe