#
#   export TFCI=...
#   export TF_BISECT_SCRIPT=ci/official/any.sh
#   export TF_BISECT_GOOD=a_good_commit_sha
#   export TF_BISECT_BAD=a_failing_commit_sha
#   export TF_ANY_TARGETS="quoted list of targets, like on the command line"
#   export TF_ANY_MODE=test
set -euxo pipefail
cd "$(dirname "$0")/../../"  # tensorflow/

# $2 in <Tensor'flow'> is <flow>. This is useful for reading string literals like below.
export TF_VER_SUFFIX=$(awk -F\" '/#define TF_VERSION_SUFFIX/ {print $2}' tensorflow/core/public/version.h)
export TF_VER_PYTHON=$(awk -F\' '/_VERSION =/ {print $2}' tensorflow/tools/pip_package/setup.py)

# Derived helper variables.
export TF_VER_SHORT="${TF_VER_MAJOR}.${TF_VER_MINOR}"

# with ample storage. When this variable is empty (i.e by default), Bazel will
# use the output base directory to run tests.
if [[ "${TFCI_MACOS_BAZEL_TEST_DIR_ENABLE}" == 1 ]]; then
  mkdir -p "${TFCI_MACOS_BAZEL_TEST_DIR_PATH}"
  export TEST_TMPDIR="${TFCI_MACOS_BAZEL_TEST_DIR_PATH}"
fi

# "TFCI_MACOS_INSTALL_BAZELISK_ENABLE" is used to decide if we need to install
# Bazelisk manually. We enable this for macOS x86 builds as those VMs do not

# Update the version numbers for Nightly only
if [[ "$TFCI_NIGHTLY_UPDATE_VERSION_ENABLE" == 1 ]]; then
  tfrun python3 tensorflow/tools/ci_build/update_version.py --nightly
  # replace tensorflow to tf_nightly in the wheel name
  export TFCI_BUILD_PIP_PACKAGE_ARGS="$(echo $TFCI_BUILD_PIP_PACKAGE_ARGS | sed 's/tensorflow/tf_nightly/')"
fi

tfrun bazel build $TFCI_BAZEL_COMMON_ARGS //tensorflow/tools/pip_package:wheel $TFCI_BUILD_PIP_PACKAGE_ARGS

# one, checks the filesize, and then ensures the new wheel is installable.
set -euxo pipefail

cd "$TFCI_OUTPUT_DIR"

# Move extra wheel files somewhere out of the way. This script
# expects just one wheel file to exist.
if [[ "$(ls *.whl | wc -l | tr -d ' ')" != "1" ]]; then
  echo "More than one wheel file is present: moving the oldest to"
  echo "$TFCI_OUTPUT_DIR/extra_wheels."

  echo 'If you have not, you will see a lot of undefined variable errors.'
else
  FROM_ENV=$(mktemp)
  # "export -p" prints a list of environment values in a safe-to-source format,
  # e.g. `declare -x TFCI_BAZEL_COMMON_ARGS="list of args"` for bash.
  export -p | grep TFCI > "$FROM_ENV"

  # Source the default ci values
  source ./ci/official/envs/ci_default

  # TODO(angerson) write this documentation

#    To use:
#       export TFCI=ci/official/envs/env_goes_here
#       export TF_ANY_TARGETS="quoted list of targets, like on the command line"
#       export TF_ANY_MODE="test" or "build" or "run" (default: "test")
#       ./any.sh
#
# 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

  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