#!/bin/bash

# Repro script for TensorFlow Python toolchain resolution problems.

set -e
set -o pipefail

# Whether to use the TF workspace rules. If not, uses an empty WORKSPACE.
tf_workspace=0
while [ "$#" -gt 0 ]; do
  case "$1" in
    --tf-workspace)
      tf_workspace=1
      shift
      ;;
    --tf-workspace-patched)
      tf_workspace_patched=1
      shift
      ;;
    *)
      printf >&2 'fatal: unknown argument "%s"' "$1"
      exit 1
      ;;
  esac
done

tmpdir="$(mktemp -d)"
workspace="${tmpdir}/src"
output_base_1="${tmpdir}/bazel_output_base_1"
output_base_2="${tmpdir}/bazel_output_base_2"
venv1="${tmpdir}/venv1"
venv2="${tmpdir}/venv2"
results="${tmpdir}/results"

# Helper to make a virtualenv and install numpy (needed for TF build).
function make-venv() {
  name="$1"
  (
    virtualenv -q --python python3 "${name}"
    source "${name}/bin/activate"
    pip -q install numpy
  )
}

# Helper to run bazel with a bunch of debugging flags.
function bazel-with-args() {
  output_base="$1"
  shift
  command="$1"
  shift
  (
  set -x
  bazel --output_base="${output_base}" "${command}" \
      --noshow_progress \
      --toolchain_resolution_debug=true \
      --experimental_repo_remote_exec \
      "$@"
  )
}

# Helper to run a single trial build.
function run_trial() {
  trial="$1"
  shift
  output_base="$1"
  shift
  venv="$1"
  shift
  (
    source "${venv}/bin/activate"
    bazel-with-args "${output_base}" build //:foo.txt
    bazel_bin="$(bazel-with-args "${output_base}" info bazel-bin)"
    cat "${bazel_bin}"/foo.txt | tee "${results}/${trial}_foo.txt"
  )
}

# Set up the example Bazel workspace.
mkdir "${workspace}"
cd "${workspace}"
touch WORKSPACE
touch tensorflow.patch

cat >optional_tensorflow.patch <<'EOF'
--- third_party/py/python_configure.bzl	1970-01-01 00:00:00.000000000 +0000
+++ third_party/py/python_configure.bzl	1970-01-01 00:00:00.000000000 +0000
@@ -268,6 +268,7 @@

 _ENVIRONS = [
     BAZEL_SH,
+    "PATH",
     PYTHON_BIN_PATH,
     PYTHON_LIB_PATH,
 ]
EOF

cat >TF_WORKSPACE <<'EOF'
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
http_archive(
    name = "io_bazel_rules_closure",
    sha256 = "6a900831c1eb8dbfc9d6879b5820fd614d4ea1db180eb5ff8aedcb75ee747c1f",
    strip_prefix = "rules_closure-db4683a2a1836ac8e265804ca5fa31852395185b",
    urls = [
        "http://mirror.tensorflow.org/github.com/bazelbuild/rules_closure/archive/db4683a2a1836ac8e265804ca5fa31852395185b.tar.gz",
        "https://github.com/bazelbuild/rules_closure/archive/db4683a2a1836ac8e265804ca5fa31852395185b.tar.gz",  # 2020-01-15
    ],
)
load("@io_bazel_rules_closure//closure:repositories.bzl", "rules_closure_dependencies")
rules_closure_dependencies(
    omit_bazel_skylib = True,
    omit_com_google_protobuf = True,
    omit_com_google_protobuf_js = True,
)
http_archive(
    name = "org_tensorflow",
    sha256 = "2595a5c401521f20a2734c4e5d54120996f8391f00bb62a57267d930bce95350",
    strip_prefix = "tensorflow-2.3.0",
    urls = [
        "http://mirror.tensorflow.org/github.com/tensorflow/tensorflow/archive/v2.3.0.tar.gz",  # 2020-07-23
        "https://github.com/tensorflow/tensorflow/archive/v2.3.0.tar.gz",
    ],
    patches = ["//:tensorflow.patch"],
)
load("@org_tensorflow//tensorflow:workspace.bzl", "tf_workspace")
tf_workspace()
EOF

cat >BUILD <<'EOF'
py_binary(
    name = "foo",
    srcs = ["foo.py"],
    python_version = "PY3",
    srcs_version = "PY3",
)

genrule(
    name = "gen_foo",
    outs = ["foo.txt"],
    cmd = "$(execpath :foo) $(execpath :foo) >$@",
    exec_tools = [":foo"],
)
EOF

cat >foo.py <<'EOF'
from __future__ import print_function
import sys
print(sys.version)
print(sys.executable)
print(sys.argv)
if len(sys.argv) >= 2:
  path_to_exec_py_binary = sys.argv[1]
  with open(path_to_exec_py_binary) as f:
    for line in f:
      if line.startswith("PYTHON_BINARY"):
        print(line.rstrip())
EOF

if [ "${tf_workspace}" -eq 1 ]; then
  # Run trials with TF workspace.
  mv TF_WORKSPACE WORKSPACE
fi
if [ "${tf_workspace_patched}" -eq 1 ]; then
  # Run trials with patched TF workspace.
  mv TF_WORKSPACE WORKSPACE
  mv optional_tensorflow.patch tensorflow.patch
fi

mkdir "${results}"

# Create the virtualenvs.
make-venv "${venv1}"
make-venv "${venv2}"

# Run trials. Switch output base without switching venv to simulate expunge;
# it's better than expunge since then we can still inspect sandbox contents.
run_trial "1-venv1" "${output_base_1}" "${venv1}"
run_trial "2-venv2" "${output_base_1}" "${venv2}"
run_trial "3-venv2-expunged" "${output_base_2}" "${venv2}"

# Print out results in consolidated format.
printf "\n\n\n"
printf "=================== RESULTS =====================\n"
for file in $(ls -1 "${results}"); do
  printf "\n===== %s =====\n" "${file}"
  cat "${results}/${file}"
done