#!/usr/bin/env bash
# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

# =============================== NOTE ===============================
# This ./pants bootstrap script comes from the pantsbuild/setup
# project. It is intended to be checked into your code repository so
# that other developers have the same setup.
#
# Learn more here: https://www.pantsbuild.org/docs/installation
# ====================================================================

set -eou pipefail

# an arbitrary number: bump when there's a change that someone might want to query for
# (e.g. checking $(PANTS_BOOTSTRAP_TOOLS=1 ./pants version) >= ...)
SCRIPT_VERSION=1

# Source any custom bootstrap settings for Pants from PANTS_BOOTSTRAP if it exists.
: ${PANTS_BOOTSTRAP:=".pants.bootstrap"}
if [[ -f "${PANTS_BOOTSTRAP}" ]]; then
  source "${PANTS_BOOTSTRAP}"
fi

# NOTE: To use an unreleased version of Pants from the pantsbuild/pants main branch,
#  locate the main branch SHA, set PANTS_SHA=<SHA> in the environment, and run this script as usual.
#
# E.g., PANTS_SHA=725fdaf504237190f6787dda3d72c39010a4c574 ./pants --version
#
# You can also use PANTS_VERSION=<VERSION> to override the config version that is in the pants.toml file.
#
# E.g., PANTS_VERSION=2.13.0 ./pants --version

PYTHON_BIN_NAME="${PYTHON:-unspecified}"

# Set this to specify a non-standard location for this script to read the Pants version from.
# NB: This will *not* cause Pants itself to use this location as a config file.
#     You can use PANTS_CONFIG_FILES or --pants-config-files to do so.
PANTS_TOML=${PANTS_TOML:-pants.toml}

PANTS_BIN_NAME="${PANTS_BIN_NAME:-$0}"

PANTS_SETUP_CACHE="${PANTS_SETUP_CACHE:-${XDG_CACHE_HOME:-$HOME/.cache}/pants/setup}"
# If given a relative path, we fix it to be absolute.
if [[ "$PANTS_SETUP_CACHE" != /* ]]; then
  PANTS_SETUP_CACHE="${PWD}/${PANTS_SETUP_CACHE}"
fi

PANTS_BOOTSTRAP="${PANTS_SETUP_CACHE}/bootstrap-$(uname -s)-$(uname -m)"

_PEX_VERSION=2.1.103
_PEX_URL="https://github.com/pantsbuild/pex/releases/download/v${_PEX_VERSION}/pex"
_PEX_EXPECTED_SHA256="4d45336511484100ae4e2bab24542a8b86b12c8cb89230463593c60d08c4b8d3"

VIRTUALENV_VERSION=20.4.7
VIRTUALENV_REQUIREMENTS=$(
cat << EOF
virtualenv==${VIRTUALENV_VERSION} --hash sha256:2b0126166ea7c9c3661f5b8e06773d28f83322de7a3ff7d06f0aed18c9de6a76
filelock==3.0.12 --hash sha256:929b7d63ec5b7d6b71b0fa5ac14e030b3f70b75747cef1b10da9b879fef15836
six==1.16.0 --hash sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
distlib==0.3.2 --hash sha256:23e223426b28491b1ced97dc3bbe183027419dfc7982b4fa2f05d5f3ff10711c
appdirs==1.4.4 --hash sha256:a841dacd6b99318a741b166adb07e19ee71a274450e68237b4650ca1055ab128
importlib-resources==5.1.4; python_version < "3.7" --hash sha256:e962bff7440364183203d179d7ae9ad90cb1f2b74dcb84300e88ecc42dca3351
importlib-metadata==4.5.0; python_version < "3.8" --hash sha256:833b26fb89d5de469b24a390e9df088d4e52e4ba33b01dc5e0e4f41b81a16c00
zipp==3.4.1; python_version < "3.10" --hash sha256:51cb66cc54621609dd593d1787f286ee42a5c0adbb4b29abea5a63edc3e03098
typing-extensions==3.10.0.0; python_version < "3.8" --hash sha256:779383f6086d90c99ae41cf0ff39aac8a7937a9283ce0a414e5dd782f4c94a84
EOF
)

COLOR_RED="\x1b[31m"
COLOR_GREEN="\x1b[32m"
COLOR_YELLOW="\x1b[33m"
COLOR_RESET="\x1b[0m"

INSTALL_URL="https://www.pantsbuild.org/docs/installation"

function log() {
  echo -e "$@" 1>&2
}

function die() {
  (($# > 0)) && log "${COLOR_RED}$*${COLOR_RESET}"
  exit 1
}

function green() {
  (($# > 0)) && log "${COLOR_GREEN}$*${COLOR_RESET}"
}

function warn() {
  (($# > 0)) && log "${COLOR_YELLOW}$*${COLOR_RESET}"
}

function tempdir {
  mkdir -p "$1"
  mktemp -d "$1"/pants.XXXXXX
}

function get_exe_path_or_die {
  local exe="$1"
  if ! command -v "${exe}"; then
    die "Could not find ${exe}. Please ensure ${exe} is on your PATH."
  fi
}

function get_pants_config_string_value {
  local config_key="$1"
  local optional_space="[[:space:]]*"
  local prefix="^${config_key}${optional_space}=${optional_space}"
  local raw_value
  raw_value="$(sed -ne "/${prefix}/ s|${prefix}||p" "${PANTS_TOML}")"
  local optional_suffix="${optional_space}(#.*)?$"
  echo "${raw_value}" \
    | sed -E \
      -e "s|^'([^']*)'${optional_suffix}|\1|" \
      -e 's|^"([^"]*)"'"${optional_suffix}"'$|\1|' \
    && return 0
  return 0
}

function get_python_major_minor_version {
  local python_exe="$1"
  "$python_exe" <<EOF
import sys
major_minor_version = ''.join(str(version_num) for version_num in sys.version_info[0:2])
print(major_minor_version)
EOF
}

# The high-level flow:
#
# 1.) Resolve the Pants version from config so that we know what interpreters we can use, what to name the venv,
#     and what to install via pip.
# 2.) Resolve the Python interpreter, first reading from the env var $PYTHON, then using a default based on the Pants
#     version.
# 3.) Check if the venv already exists via a naming convention, and create the venv if not found.
# 4.) Execute Pants with the resolved Python interpreter and venv.
#
# After that, Pants itself will handle making sure any requested plugins
# are installed and up to date.

function determine_pants_version {
  if [ -n "${PANTS_SHA:-}" ]; then
    # get_version_for_sha will echo the version, thus "returning" it from this function.
    get_version_for_sha "$PANTS_SHA"
    return
  fi

  if [ -n "${PANTS_VERSION:-}" ]; then
	echo "${PANTS_VERSION}"
    return
  fi

  pants_version="$(get_pants_config_string_value 'pants_version')"
  if [[ -z "${pants_version}" ]]; then
    die "Please explicitly specify the \`pants_version\` in your \`pants.toml\` under the \`[GLOBAL]\` scope.
See https://pypi.org/project/pantsbuild.pants/#history for all released versions
and ${INSTALL_URL} for more instructions."
  fi
  pants_major_version="$(echo "${pants_version}" | cut -d '.' -f1)"
  pants_minor_version="$(echo "${pants_version}" | cut -d '.' -f2)"
  # 1.26 is the first version to support `pants.toml`, so we fail eagerly if using an outdated version.
  if [[ "${pants_major_version}" -eq 1 && "${pants_minor_version}" -le 25 ]]; then
    die "This version of the \`./pants\` script does not work with Pants <= 1.25.0 (and it also requires using \`pants.toml\`,
rather than \`pants.ini\`). Instead, either upgrade your \`pants_version\` or use the version of the \`./pants\` script
at https://raw.githubusercontent.com/Eric-Arellano/setup/0d445edef57cb89fd830db70810e38f050b0a268/pants."
  fi
  echo "${pants_version}"
}

function set_supported_python_versions {
  local pants_version="$1"
  local pants_major_version
  local pants_minor_version
  pants_major_version="$(echo "${pants_version}" | cut -d '.' -f1)"
  pants_minor_version="$(echo "${pants_version}" | cut -d '.' -f2)"
  if [[ "${pants_major_version}" -eq 1 ]]; then
    supported_python_versions_decimal=('3.6' '3.7' '3.8')
    supported_python_versions_int=('36' '37' '38')
    supported_message='3.6, 3.7, or 3.8'
  elif [[ "${pants_major_version}" -eq 2 && "${pants_minor_version}" -eq 0 ]]; then
    supported_python_versions_decimal=('3.6' '3.7' '3.8')
    supported_python_versions_int=('36' '37' '38')
    supported_message='3.6, 3.7, or 3.8'
  elif [[ "${pants_major_version}" -eq 2 && "${pants_minor_version}" -eq 1 ]]; then
    supported_python_versions_decimal=('3.7' '3.8' '3.6')
    supported_python_versions_int=('37' '38' '36')
    supported_message='3.7, 3.8, or 3.6 (deprecated)'
  elif [[ "${pants_major_version}" -eq 2 && "${pants_minor_version}" -lt 5 ]]; then
    supported_python_versions_decimal=('3.8' '3.7')
    supported_python_versions_int=('38' '37')
    supported_message='3.7 or 3.8'
  else
    # We put 3.9 first because Apple Silicon only works properly with Python 3.9, even though it's possible to have
    # older Pythons installed. This makes it more likely that Pants will work out-of-the-box.
    supported_python_versions_decimal=('3.9' '3.8' '3.7')
    supported_python_versions_int=('39' '38' '37')
    supported_message='3.7, 3.8, or 3.9'
  fi
}

function check_python_exe_compatible_version {
  local python_exe="$1"
  local major_minor_version
  major_minor_version="$(get_python_major_minor_version "${python_exe}")"
  for valid_version in "${supported_python_versions_int[@]}"; do
    if [[ "${major_minor_version}" == "${valid_version}" ]]; then
      echo "${python_exe}" && return 0
    fi
  done
}

function determine_default_python_exe {
  for version in "${supported_python_versions_decimal[@]}" "3" ""; do
    local interpreter_path
    interpreter_path="$(command -v "python${version}")"
    if [[ -z "${interpreter_path}" ]]; then
      continue
    fi
    # Check if a version is shimmed by pyenv or asdf but not configured.
    if ! "${interpreter_path}" --version >/dev/null 2>&1; then
      continue
    fi
    if [[ -n "$(check_python_exe_compatible_version "${interpreter_path}")" ]]; then
      echo "${interpreter_path}" && return 0
    fi
  done
}

function determine_python_exe {
  local pants_version="$1"
  set_supported_python_versions "${pants_version}"
  local requirement_str="For \`pants_version = \"${pants_version}\"\`, Pants requires Python ${supported_message} to run."

  local python_exe
  if [[ "${PYTHON_BIN_NAME}" != 'unspecified' ]]; then
    python_exe="$(get_exe_path_or_die "${PYTHON_BIN_NAME}")" || exit 1
    if [[ -z "$(check_python_exe_compatible_version "${python_exe}")" ]]; then
      die "Invalid Python interpreter version for ${python_exe}. ${requirement_str}"
    fi
  else
    python_exe="$(determine_default_python_exe)"
    if [[ -z "${python_exe}" ]]; then
      die "No valid Python interpreter found. ${requirement_str} Please check that a valid interpreter is installed and on your \$PATH."
    fi
  fi
  echo "${python_exe}"
}

function compute_sha256 {
  local python="$1"
  local path="$2"

  "$python" <<EOF
import hashlib

hasher = hashlib.sha256()
with open('${path}', 'rb') as fp:
    buf = fp.read()
    hasher.update(buf)
print(hasher.hexdigest())
EOF
}

# TODO(John Sirois): GC race loser tmp dirs leftover from bootstrap_XXX
# functions.  Any tmp dir w/o a symlink pointing to it can go.

function bootstrap_pex {
  local python="$1"
  local bootstrapped="${PANTS_BOOTSTRAP}/pex-${_PEX_VERSION}/pex"
  if [[ ! -f "${bootstrapped}" ]]; then
    (
      green "Downloading the Pex PEX."
      mkdir -p "${PANTS_BOOTSTRAP}"
      local staging_dir
      staging_dir=$(tempdir "${PANTS_BOOTSTRAP}")
      curl --proto "=https" \
           --tlsv1.2 \
           --silent \
           --location \
           -o "${staging_dir}/pex" \
           "${_PEX_URL}"
      fingerprint="$(compute_sha256 "${python}" "${staging_dir}/pex")"
      if [[ "${_PEX_EXPECTED_SHA256}" != "${fingerprint}" ]]; then
        die "SHA256 of ${_PEX_URL} is not as expected. Aborting."
      fi
      green "SHA256 fingerprint of ${_PEX_URL} verified."
      mkdir -p "$(dirname "${bootstrapped}")"
      mv -f "${staging_dir}/pex" "${bootstrapped}"
      rmdir "${staging_dir}"
    ) 1>&2 || exit 1
  fi
  echo "${bootstrapped}"
}

function scrub_env_vars {
  # Ensure the virtualenv PEX runs as shrink-wrapped.
  # See: https://github.com/pantsbuild/setup/issues/105
  local -r pex_env_vars=(${!PEX_@})
  if [[ ! ${#pex_env_vars[@]} -eq 0 ]]; then
    local -r pex_env_vars_to_scrub="${pex_env_vars[@]/PEX_ROOT}"
    if [[ -n "${pex_env_vars_to_scrub[@]}" ]]; then
      warn "Scrubbing ${pex_env_vars_to_scrub[@]}"
      unset ${pex_env_vars_to_scrub[@]}
    fi
  fi
  # Also ensure pip doesn't think packages on PYTHONPATH
  # are already installed.
  if [ -n "${PYTHONPATH:-}" ]; then
    warn "Scrubbing PYTHONPATH"
    unset PYTHONPATH
  fi
}

function bootstrap_virtualenv {
  local python="$1"
  local bootstrapped="${PANTS_BOOTSTRAP}/virtualenv-${VIRTUALENV_VERSION}/virtualenv.pex"
  if [[ ! -f "${bootstrapped}" ]]; then
    (
      green "Creating the virtualenv PEX."
      pex_path="$(bootstrap_pex "${python}")" || exit 1
      mkdir -p "${PANTS_BOOTSTRAP}"
      local staging_dir
      staging_dir=$(tempdir "${PANTS_BOOTSTRAP}")
      echo "${VIRTUALENV_REQUIREMENTS}" > "${staging_dir}/requirements.txt"
      (
        scrub_env_vars
        "${python}" "${pex_path}" -r "${staging_dir}/requirements.txt" -c virtualenv -o "${staging_dir}/virtualenv.pex"
      )
      mkdir -p "$(dirname "${bootstrapped}")"
      mv -f "${staging_dir}/virtualenv.pex" "${bootstrapped}"
      rm -rf "${staging_dir}"
    ) 1>&2 || exit 1
  fi
  echo "${bootstrapped}"
}

function find_links_url {
  local pants_version="$1"
  local pants_sha="$2"
  echo -n "https://binaries.pantsbuild.org/wheels/pantsbuild.pants/${pants_sha}/${pants_version/+/%2B}/index.html"
}

function get_version_for_sha {
  local sha="$1"

  # Retrieve the Pants version associated with this commit.
  local pants_version
  pants_version="$(curl --proto "=https" \
                        --tlsv1.2 \
                        --fail \
                        --silent \
                        --location \
                        "https://raw.githubusercontent.com/pantsbuild/pants/${sha}/src/python/pants/VERSION")"

  # Construct the version as the release version from src/python/pants/VERSION, plus the string `+gitXXXXXXXX`,
  # where the XXXXXXXX is the first 8 characters of the SHA.
  echo "${pants_version}+git${sha:0:8}"
}

function bootstrap_pants {
  local pants_version="$1"
  local python="$2"
  local pants_sha="${3:-}"
  local pants_debug="${4:-}"

  local pants_requirements=(pantsbuild.pants==${pants_version})
  local maybe_find_links
  if [[ -z "${pants_sha}" ]]; then
    maybe_find_links=""
  else
    maybe_find_links="--find-links=$(find_links_url "${pants_version}" "${pants_sha}")"
  fi

  local debug_suffix
  if [[ -z "${pants_debug}" ]]; then
    debug_suffix=""
  else
    debug_suffix="-debug"
    pants_requirements+=(debugpy==1.6.0)
  fi

  local python_major_minor_version
  python_major_minor_version="$(get_python_major_minor_version "${python}")"
  local target_folder_name="${pants_version}_py${python_major_minor_version}${debug_suffix}"
  local bootstrapped="${PANTS_BOOTSTRAP}/${target_folder_name}"

  if [[ ! -d "${bootstrapped}" ]]; then
    (
      green "Bootstrapping Pants using ${python}"
      local staging_dir
      staging_dir=$(tempdir "${PANTS_BOOTSTRAP}")
      local virtualenv_path
      virtualenv_path="$(bootstrap_virtualenv "${python}")" || exit 1
      green "Installing ${pants_requirements[@]} into a virtual environment at ${bootstrapped}"
      (
        scrub_env_vars
        # shellcheck disable=SC2086
        "${python}" "${virtualenv_path}" --quiet --no-download "${staging_dir}/install" && \
        # Grab the latest pip, but don't advance setuptools past 58 which drops support for the
        # `setup` kwarg `use_2to3` which Pants 1.x sdist dependencies (pystache) use.
        "${staging_dir}/install/bin/pip" install --quiet -U pip "setuptools<58" && \
        "${staging_dir}/install/bin/pip" install ${maybe_find_links} --quiet --progress-bar off "${pants_requirements[@]}"
      ) && \
      ln -s "${staging_dir}/install" "${staging_dir}/${target_folder_name}" && \
      mv "${staging_dir}/${target_folder_name}" "${bootstrapped}" && \
      green "New virtual environment successfully created at ${bootstrapped}."
    ) 1>&2 || exit 1
  fi
  echo "${bootstrapped}"
}

function run_bootstrap_tools {
  # functionality for introspecting the bootstrapping process, without actually doing it
  if [[ "${PANTS_BOOTSTRAP_TOOLS}" -gt "${SCRIPT_VERSION}" ]]; then
      die "$0 script (bootstrap version ${SCRIPT_VERSION}) is too old for this invocation (with PANTS_BOOTSTRAP_TOOLS=${PANTS_BOOTSTRAP_TOOLS}).
Please update it by following ${INSTALL_URL}"
  fi

  case "${1:-}" in
    bootstrap-cache-key)
      local pants_version=$(determine_pants_version)
      local python="$(determine_python_exe "${pants_version}")"
      # the python above may be a shim (e.g. pyenv or homebrew), so let's get an estimate of the
      # actual path, as will be symlinked in the virtualenv. (NB. virtualenv does more complicated
      # things, but we at least emulate the symlink-resolution that it does.)
      local python_executable_path="$("${python}" -c 'import os, sys; print(os.path.realpath(sys.executable))')"

      local requirements_file="$(mktemp)"
      echo "${VIRTUALENV_REQUIREMENTS}" > "${requirements_file}"
      local virtualenv_requirements_sha256="$(compute_sha256 "${python}" "${requirements_file}")"
      rm "${requirements_file}"

      local parts=(
        "os_name=$(uname -s)"
        "arch=$(uname -m)"
        "python_path=${python}"
        "python_executable_path=${python_executable_path}"
        # the full interpreter information, for maximum compatibility
        "python_version=$("$python" --version)"
        "pex_version=${_PEX_VERSION}"
        "virtualenv_requirements_sha256=${virtualenv_requirements_sha256}"
        "pants_version=${pants_version}"
      )
      echo "${parts[*]}"
      ;;
    bootstrap-version)
      echo "${SCRIPT_VERSION}"
      ;;
    help|"")
      cat <<EOF
Usage: PANTS_BOOTSTRAP_TOOLS=1 $0 ...

Subcommands:
  bootstrap-cache-key
    Print an opaque that can be used as a key for accurate and safe caching of
    the pants bootstrap directories.

    (Added in bootstrap version 1.)

  bootstrap-version
    Print a version number for the bootstrap script itself.

    Distributed scripts (such as reusable CI formulae) that use these bootstrap
    tools should set PANTS_BOOTSTRAP_TOOLS to the minimum script version for the
    features they require. For example, if 'some-tool' was added in version 123:

        PANTS_BOOTSTRAP_TOOLS=123 ./pants some-tool

    (Added in bootstrap version 1.)
EOF
      ;;
    *)
      die "Unknown subcommand for bootstrap tools: $1. Do you mean to run without PANTS_BOOTSTRAP_TOOLS=1? Or, update this script ($INSTALL_URL)?"
  esac
}

if [[ "${PANTS_BOOTSTRAP_TOOLS:-}" -gt 0 ]]; then
    run_bootstrap_tools "$@"
    exit 0
fi

# Ensure we operate from the context of the ./pants buildroot.
cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"
pants_version="$(determine_pants_version)"
python="$(determine_python_exe "${pants_version}")"
pants_dir="$(bootstrap_pants "${pants_version}" "${python}" "${PANTS_SHA:-}" "${PANTS_DEBUG:-}")" || exit 1

pants_python="${pants_dir}/bin/python"
pants_binary=(${pants_dir}/bin/pants)
pants_extra_args=""
if [[ -n "${PANTS_SHA:-}" ]]; then
  pants_extra_args="${pants_extra_args} --python-repos-repos=$(find_links_url "$pants_version" "$PANTS_SHA")"
fi

pants_debug_args=()
if [ -n "${PANTS_DEBUG:-}" ]; then
  if [[ "$*" != *"--no-pantsd"* ]]; then
    echo "Error! Must pass '--no-pantsd' when using PANTS_DEBUG"
    exit 1
  fi
  # NB: We can't invoke `-m debugpy` as that'll prepend CWD to sys.path which might have unintended side-effects.
  # `-c` also prepends, but we strip that ourselves.
  pants_binary=(-c "__import__(\"sys\").path.pop(0);__import__(\"debugpy.server.cli\").server.cli.main()" --listen 127.0.0.1:5678 --wait-for-client "${pants_binary[@]}")
  echo "Will launch debugpy server at '127.0.0.1:5678' waiting for client connection."
fi

# shellcheck disable=SC2086
exec "${pants_python}" "${pants_binary[@]}" ${pants_extra_args} \
  --pants-bin-name="${PANTS_BIN_NAME}" --pants-version=${pants_version} "$@"