# Copyright 2014 Pex project contributors.
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import absolute_import, print_function

import contextlib
import glob
import itertools
import os
import platform
import random
import re
import sys
from collections import Counter
from contextlib import contextmanager
from textwrap import dedent

import pytest

from pex.atomic_directory import atomic_directory
from pex.common import open_zip, safe_mkdir, safe_mkdtemp, safe_rmtree, safe_sleep, temporary_dir
from pex.compatibility import to_unicode
from pex.dist_metadata import Distribution
from pex.executor import Executor
from pex.interpreter import PythonInterpreter
from pex.interpreter_implementation import InterpreterImplementation
from pex.os import LINUX, MAC, WINDOWS
from pex.pep_427 import install_wheel_chroot, install_wheel_interpreter
from pex.pex import PEX
from pex.pex_builder import PEXBuilder
from pex.pex_info import PexInfo
from pex.pip.installation import get_pip
from pex.pip.version import PipVersion, PipVersionValue
from pex.resolve.configured_resolver import ConfiguredResolver
from pex.sysconfig import SCRIPT_DIR, script_name
from pex.targets import LocalInterpreter
from pex.typing import TYPE_CHECKING, cast
from pex.util import CacheHelper, named_temporary_file
from pex.venv.virtualenv import InstallationChoice, Virtualenv

# Explicitly re-export subprocess to enable a transparent substitution in tests that supports
# executing PEX files directly on Windows.
from testing import subprocess as subprocess
from testing.pytest_utils.tmp import Tempdir

try:
    from unittest import mock
except ImportError:
    import mock  # type: ignore[no-redef,import]

if TYPE_CHECKING:
    from typing import (
        Any,
        Callable,
        Dict,
        Iterable,
        Iterator,
        List,
        Mapping,
        Optional,
        Set,
        Text,
        Tuple,
        Union,
    )

    import attr  # vendor:skip
else:
    from pex.third_party import attr

PY_VER = sys.version_info[:2]
IS_PYPY = hasattr(sys, "pypy_version_info")
IS_PYPY2 = IS_PYPY and sys.version_info[0] == 2
IS_PYPY3 = IS_PYPY and sys.version_info[0] == 3
NOT_CPYTHON27 = IS_PYPY or PY_VER != (2, 7)
IS_LINUX = LINUX
IS_MAC = MAC
IS_WINDOWS = WINDOWS
IS_X86_64 = platform.machine().lower() in ("amd64", "x86_64")
IS_ARM_64 = platform.machine().lower() in ("arm64", "aarch64")
IS_LINUX_X86_64 = IS_LINUX and IS_X86_64
IS_LINUX_ARM64 = IS_LINUX and IS_ARM_64
IS_MAC_X86_64 = IS_MAC and IS_X86_64
IS_MAC_ARM64 = IS_MAC and IS_ARM_64
NOT_CPYTHON27_OR_OSX = NOT_CPYTHON27 or not IS_LINUX


@contextlib.contextmanager
def temporary_filename():
    # type: () -> Iterator[str]
    """Creates a temporary filename.

    This is useful when you need to pass a filename to an API. Windows requires all handles to a
    file be closed before deleting/renaming it, so this makes it a bit simpler.
    """
    with named_temporary_file() as fp:
        fp.write(b"")
        fp.close()
        yield fp.name


def random_bytes(length):
    # type: (int) -> bytes
    return "".join(map(chr, (random.randint(ord("a"), ord("z")) for _ in range(length)))).encode(
        "utf-8"
    )


def get_dep_dist_names_from_pex(pex_path, match_prefix=""):
    # type: (str, str) -> Set[Text]
    """Given an on-disk pex, extract all of the unique first-level paths under `.deps`."""
    with open_zip(pex_path) as pex_zip:
        dep_gen = (f.split(os.sep)[1] for f in pex_zip.namelist() if f.startswith(".deps/"))
        return set(item for item in dep_gen if item.startswith(match_prefix))


@contextlib.contextmanager
def temporary_content(content_map, interp=None, seed=31337, perms=0o644):
    # type: (Mapping[str, Union[int, str]], Optional[Dict[str, Any]], int, int) -> Iterator[str]
    """Write content to disk where content is map from string => (int, string).

    If target is int, write int random bytes.  Otherwise write contents of string.
    """
    random.seed(seed)
    interp = interp or {}
    with temporary_dir() as td:
        for filename, size_or_content in content_map.items():
            dest = os.path.join(td, filename)
            safe_mkdir(os.path.dirname(dest))
            with open(dest, "wb") as fp:
                if isinstance(size_or_content, int):
                    fp.write(random_bytes(size_or_content))
                else:
                    fp.write((size_or_content % interp).encode("utf-8"))
            os.chmod(dest, perms)
        yield td


@contextlib.contextmanager
def make_project(
    name="my_project",  # type: str
    version="0.0.0",  # type: str
    zip_safe=True,  # type: bool
    install_reqs=None,  # type: Optional[List[str]]
    extras_require=None,  # type: Optional[Dict[str, List[str]]]
    entry_points=None,  # type: Optional[Union[str, Dict[str, List[str]]]]
    python_requires=None,  # type: Optional[str]
    universal=False,  # type: bool
    prepare_project=None,  # type: Optional[Callable[[str], None]]
):
    # type: (...) -> Iterator[str]
    project_content = {
        "setup.py": dedent(
            """
            from setuptools import setup
            
            setup(
            name=%(project_name)r,
            version=%(version)r,
            zip_safe=%(zip_safe)r,
            packages=[%(project_name)r],
            scripts=[
              'scripts/%(hello_world_script_name)s',
              'scripts/%(shell_script_name)s',
            ],
            package_data={%(project_name)r: ['package_data/*.dat']},
            install_requires=%(install_requires)r,
            extras_require=%(extras_require)r,
            entry_points=%(entry_points)r,
            python_requires=%(python_requires)r,
            options={'bdist_wheel': {'universal': %(universal)r}},
            )
            """
        ),
        os.path.join(name, "__init__.py"): 0,
        os.path.join(name, "my_module.py"): 'def do_something():\n  print("hello world!")\n',
        os.path.join(name, "package_data/resource1.dat"): 1000,
        os.path.join(name, "package_data/resource2.dat"): 1000,
    }  # type: Dict[str, Union[str, int]]

    if WINDOWS:
        project_content.update(
            (
                (
                    "scripts/hello_world.py",
                    '#!/usr/bin/env python\r\nprint("hello world from py script!")\r\n',
                ),
                ("scripts/shell_script.bat", "@echo off\r\necho hello world from shell script\r\n"),
            )
        )
    else:
        project_content.update(
            (
                (
                    "scripts/hello_world",
                    '#!/usr/bin/env python\nprint("hello world from py script!")\n',
                ),
                (
                    "scripts/shell_script",
                    "#!/usr/bin/env bash\necho hello world from shell script\n",
                ),
            )
        )

    interp = {
        "project_name": name,
        "version": version,
        "zip_safe": zip_safe,
        "install_requires": install_reqs or [],
        "extras_require": extras_require or {},
        "entry_points": entry_points or {},
        "python_requires": python_requires,
        "universal": universal,
        "hello_world_script_name": "hello_world.py" if WINDOWS else "hello_world",
        "shell_script_name": "shell_script.bat" if WINDOWS else "shell_script",
    }

    with temporary_content(project_content, interp=interp) as td:
        if prepare_project:
            prepare_project(td)
        yield td


class WheelBuilder(object):
    """Create a wheel distribution from an unpacked setup.py-based project."""

    class BuildFailure(Exception):
        pass

    def __init__(
        self,
        source_dir,  # type: str
        interpreter=None,  # type: Optional[PythonInterpreter]
        pip_version=None,  # type: Optional[PipVersionValue]
        wheel_dir=None,  # type: Optional[str]
        verify=True,  # type: bool
    ):
        # type: (...) -> None
        """Create a wheel from an unpacked source distribution in source_dir."""
        self._source_dir = source_dir
        self._wheel_dir = wheel_dir or safe_mkdtemp()
        self._interpreter = interpreter or PythonInterpreter.get()
        self._pip_version = pip_version or PipVersion.DEFAULT
        self._resolver = ConfiguredResolver.version(self._pip_version)
        self._verify = verify

    def bdist(self):
        # type: () -> str
        get_pip(
            interpreter=self._interpreter,
            version=self._pip_version,
            resolver=self._resolver,
        ).spawn_build_wheels(
            requirements=[self._source_dir],
            wheel_dir=self._wheel_dir,
            interpreter=self._interpreter,
            verify=self._verify,
        ).wait()
        dists = glob.glob(os.path.join(self._wheel_dir, "*.whl"))
        if len(dists) == 0:
            raise self.BuildFailure("No distributions were produced!")
        if len(dists) > 1:
            raise self.BuildFailure(
                "Ambiguous wheel distributions found: {dists}".format(dists=" ".join(dists))
            )
        return dists[0]


@contextlib.contextmanager
def built_wheel(
    name="my_project",  # type: str
    version="0.0.0",  # type: str
    zip_safe=True,  # type: bool
    install_reqs=None,  # type: Optional[List[str]]
    extras_require=None,  # type: Optional[Dict[str, List[str]]]
    entry_points=None,  # type: Optional[Union[str, Dict[str, List[str]]]]
    interpreter=None,  # type: Optional[PythonInterpreter]
    python_requires=None,  # type: Optional[str]
    universal=False,  # type: bool
    prepare_project=None,  # type: Optional[Callable[[str], None]]
    **kwargs  # type: Any
):
    # type: (...) -> Iterator[str]
    with make_project(
        name=name,
        version=version,
        zip_safe=zip_safe,
        install_reqs=install_reqs,
        extras_require=extras_require,
        entry_points=entry_points,
        python_requires=python_requires,
        universal=universal,
        prepare_project=prepare_project,
    ) as td:
        builder = WheelBuilder(td, interpreter=interpreter, **kwargs)
        yield builder.bdist()


@contextlib.contextmanager
def make_source_dir(
    name="my_project",  # type: str
    version="0.0.0",  # type: str
    install_reqs=None,  # type: Optional[List[str]]
    extras_require=None,  # type: Optional[Dict[str, List[str]]]
):
    # type: (...) -> Iterator[str]
    with make_project(
        name=name, version=version, install_reqs=install_reqs, extras_require=extras_require
    ) as td:
        yield td


@contextlib.contextmanager
def make_bdist(
    name="my_project",  # type: str
    version="0.0.0",  # type: str
    zip_safe=True,  # type: bool
    interpreter=None,  # type: Optional[PythonInterpreter]
    **kwargs  # type: Any
):
    # type: (...) -> Iterator[Distribution]
    with built_wheel(
        name=name, version=version, zip_safe=zip_safe, interpreter=interpreter, **kwargs
    ) as dist_location:
        yield install_wheel(dist_location)


def install_wheel(wheel):
    # type: (str) -> Distribution
    install_dir = os.path.join(safe_mkdtemp(), os.path.basename(wheel))
    install_wheel_chroot(wheel=wheel, destination=install_dir)
    return Distribution.load(install_dir)


COVERAGE_PREAMBLE = """
try:
  from coverage import coverage
  cov = coverage(auto_data=True, data_suffix=True)
  cov.start()
except ImportError:
  pass
"""


def write_simple_pex(
    td,  # type: str
    exe_contents=None,  # type: Optional[str]
    dists=None,  # type: Optional[Iterable[Distribution]]
    sources=None,  # type: Optional[Iterable[Tuple[str, str]]]
    coverage=False,  # type: bool
    interpreter=None,  # type: Optional[PythonInterpreter]
    pex_info=None,  # type: Optional[PexInfo]
):
    # type: (...) -> PEXBuilder
    """Write a pex file that optionally contains an executable entry point.

    :param td: temporary directory path
    :param exe_contents: entry point python file
    :param dists: distributions to include, typically sdists or bdists
    :param sources: sources to include, as a list of pairs (env_filename, contents)
    :param coverage: include coverage header
    :param interpreter: a custom interpreter to use to build the pex
    :param pex_info: a custom PexInfo to use to build the pex.
    """
    dists = dists or []
    sources = sources or []

    safe_mkdir(td)

    pb = PEXBuilder(
        path=td,
        preamble=COVERAGE_PREAMBLE if coverage else None,
        interpreter=interpreter,
        pex_info=pex_info,
    )

    for dist in dists:
        pb.add_dist_location(dist.location if isinstance(dist, Distribution) else dist)

    for env_filename, contents in sources:
        src_path = os.path.join(td, env_filename)
        safe_mkdir(os.path.dirname(src_path))
        with open(src_path, "w") as fp:
            fp.write(contents)
        pb.add_source(src_path, env_filename)

    if exe_contents:
        with open(os.path.join(td, "exe.py"), "w") as fp:
            fp.write(exe_contents)
        pb.set_executable(os.path.join(td, "exe.py"))

    pb.freeze()

    return pb


def re_exact(text):
    # type: (str) -> str
    return r"^{escaped}$".format(escaped=re.escape(text))


@attr.s(frozen=True)
class IntegResults(object):
    """Convenience object to return integration run results."""

    cmd = attr.ib()  # type: Tuple[str, ...]
    output = attr.ib()  # type: Text
    error = attr.ib()  # type: Text
    return_code = attr.ib()  # type: int

    @property
    def exe(self):
        # type: () -> str
        return self.cmd[0]

    @property
    def interpreter(self):
        # type: () -> PythonInterpreter
        return PythonInterpreter.from_binary(self.exe)

    @property
    def target(self):
        # type: () -> LocalInterpreter
        return LocalInterpreter.create(self.exe)

    @property
    def pex(self):
        # type: () -> PEX
        return PEX(self.exe)

    def assert_success(
        self,
        expected_output_re=None,  # type: Optional[str]
        expected_error_re=None,  # type: Optional[str]
        re_flags=0,  # type: int
    ):
        # type: (...) -> None
        assert self.return_code == 0, to_unicode(
            "integration test failed: return_code={return_code}, output={output}, error={error}"
        ).format(return_code=self.return_code, output=self.output, error=self.error)
        self.assert_output(expected_output_re, expected_error_re, re_flags)

    def assert_failure(
        self,
        expected_error_re=None,  # type: Optional[str]
        expected_output_re=None,  # type: Optional[str]
        re_flags=0,  # type: int
    ):
        # type: (...) -> None
        assert self.return_code != 0
        self.assert_output(expected_output_re, expected_error_re, re_flags)

    def assert_output(
        self,
        expected_output_re=None,  # type: Optional[str]
        expected_error_re=None,  # type: Optional[str]
        re_flags=0,  # type: int
    ):
        if expected_output_re:
            assert re.match(expected_output_re, self.output, flags=re_flags), to_unicode(
                "Failed to match re: {re!r} against:\n{output}"
            ).format(re=expected_output_re, output=self.output)
        if expected_error_re:
            assert re.match(expected_error_re, self.error, flags=re_flags), to_unicode(
                "Failed to match re: {re!r} against:\n{output}"
            ).format(re=expected_error_re, output=self.error)


def create_pex_command(
    args=None,  # type: Optional[Iterable[str]]
    python=None,  # type: Optional[str]
    quiet=None,  # type: Optional[bool]
    pex_module="pex",  # type: str
    use_pex_whl_venv=True,  # type: bool
):
    # type: (...) -> List[str]
    python_exe = python or sys.executable
    cmd = [
        installed_pex_wheel_venv_python(python_exe) if use_pex_whl_venv else python_exe,
        "-m",
        pex_module,
    ]
    if pex_module == "pex" and not quiet:
        cmd.append("-v")
    if args:
        cmd.extend(args)
    return cmd


@attr.s(frozen=True)
class UvPython(object):
    @staticmethod
    def _get_os():
        # type: () -> str

        if IS_WINDOWS:
            return "windows"
        elif IS_MAC:
            return "macos"
        else:
            return "linux"

    @staticmethod
    def _get_arch():
        # type: () -> str

        if IS_WINDOWS:
            arch = os.environ.get("PROCESSOR_ARCHITECTURE", platform.machine()).lower()
        else:
            arch = platform.machine().lower()

        if arch in ("aarch64", "arm64"):
            return "aarch64"
        elif arch in ("x86_64", "amd64"):
            return "x86_64"
        else:
            return arch

    major = attr.ib()  # type: int
    minor = attr.ib()  # type: int
    implementation = attr.ib(
        default=InterpreterImplementation.CPYTHON
    )  # type: InterpreterImplementation.Value

    def spec(self):
        # type: () -> str

        if self.implementation is InterpreterImplementation.PYPY:
            return "pypy-{major}.{minor}-{os}-{arch}".format(
                major=self.major, minor=self.minor, os=self._get_os(), arch=self._get_arch()
            )

        # N.B.: We force arch to get arm64 PBS builds for Windows arm64 machines.
        # See: https://github.com/astral-sh/uv/issues/12906
        return "cpython-{major}.{minor}{suffix}-{os}-{arch}".format(
            major=self.major,
            minor=self.minor,
            suffix="t"
            if self.implementation is InterpreterImplementation.CPYTHON_FREE_THREADED
            else "",
            os=self._get_os(),
            arch=self._get_arch(),
        )


def ensure_uv_python(
    python,  # type: Union[UvPython, str]
    install_if_missing=True,  # type: bool
):
    # type: (...) -> str

    python_spec = python.spec() if isinstance(python, UvPython) else python

    env = make_env(UV_PYTHON_INSTALL_DIR=os.path.join(PEX_TEST_DEV_ROOT, "uv-pythons"))
    process = subprocess.Popen(
        args=["uv", "python", "find", python_spec], env=env, stdout=subprocess.PIPE
    )
    stdout, _ = process.communicate()
    if process.returncode == 0:
        return cast(str, stdout.decode("utf-8").strip())

    assert install_if_missing, "The {python} interpreter was not found by uv on the system.".format(
        python=python
    )
    subprocess.check_call(
        args=["uv", "python", "install", "--managed-python", "--force", python_spec], env=env
    )
    return ensure_uv_python(python_spec, install_if_missing=False)


def run_pex_command(
    args,  # type: Iterable[str]
    env=None,  # type: Optional[Dict[str, str]]
    python=None,  # type: Optional[Union[str, UvPython]]
    quiet=None,  # type: Optional[bool]
    cwd=None,  # type: Optional[str]
    pex_module="pex",  # type: str
    use_pex_whl_venv=True,  # type: bool
):
    # type: (...) -> IntegResults
    """Simulate running pex command for integration testing.

    This is different from run_simple_pex in that it calls the pex command rather than running a
    generated pex.  This is useful for testing end to end runs with specific command line arguments
    or env options.
    """
    if isinstance(python, UvPython):
        python = ensure_uv_python(python)
    cmd = create_pex_command(
        args, python=python, quiet=quiet, pex_module=pex_module, use_pex_whl_venv=use_pex_whl_venv
    )
    process = Executor.open_process(
        cmd=cmd, env=env, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    output, error = process.communicate()
    return IntegResults(
        tuple(cmd), output.decode("utf-8"), error.decode("utf-8"), process.returncode
    )


def run_simple_pex(
    pex,  # type: str
    args=(),  # type: Iterable[str]
    interpreter=None,  # type: Optional[PythonInterpreter]
    stdin=None,  # type: Optional[bytes]
    **kwargs  # type: Any
):
    # type: (...) -> Tuple[bytes, int]
    p = PEX(pex, interpreter=interpreter)
    process = p.run(
        args=args,
        blocking=False,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=kwargs.pop("stderr", subprocess.STDOUT),
        **kwargs
    )
    stdout, _ = process.communicate(input=stdin)
    return stdout.replace(b"\r", b""), process.returncode


def run_simple_pex_test(
    body,  # type: str
    args=(),  # type: Iterable[str]
    env=None,  # type: Optional[Mapping[str, str]]
    dists=None,  # type: Optional[Iterable[Distribution]]
    coverage=False,  # type: bool
    interpreter=None,  # type: Optional[PythonInterpreter]
):
    # type: (...) -> Tuple[bytes, int]
    with temporary_dir() as td1, temporary_dir() as td2:
        pb = write_simple_pex(td1, body, dists=dists, coverage=coverage, interpreter=interpreter)
        pex = os.path.join(td2, "app.pex")
        pb.build(pex)
        return run_simple_pex(pex, args=args, env=env, interpreter=interpreter)


PYENV_GIT_URL = "https://github.com/{pyenv}".format(
    pyenv="pyenv-win/pyenv-win" if WINDOWS else "pyenv/pyenv"
)


def bootstrap_python_installer(dest):
    # type: (str) -> None
    for index in range(3):
        try:
            subprocess.check_call(args=["git", "clone", "--depth", "1", PYENV_GIT_URL, dest])
            return
        except subprocess.CalledProcessError as e:
            print("Error cloning pyenv on attempt", index + 1, "of 3:", e, file=sys.stderr)
            continue
    raise RuntimeError("Could not clone pyenv from git after 3 tries.")


# NB: We keep the pool of bootstrapped interpreters as small as possible to avoid timeouts in CI
# otherwise encountered when fetching and building too many on a cache miss. In the past we had
# issues with the combination of 7 total unique interpreter versions and a Travis-CI timeout of 50
# minutes for a shard.
# N.B.: Make sure to stick to versions that have binary releases for all supported platforms to
# support use of pyenv-win which does not build from source, just running released installers
# instead.
PY27 = "2.7.18"
PY39 = "3.9.13"
PY310 = "3.10.7"
PY311 = "3.11.13"

ALL_PY_VERSIONS = (PY27, PY39, PY310, PY311)
_ALL_PY_VERSIONS_TO_VERSION_INFO = {
    version: tuple(map(int, version.split("."))) for version in ALL_PY_VERSIONS
}


PEX_TEST_DEV_ROOT = os.path.abspath(
    os.path.expanduser(os.environ.get("_PEX_TEST_DEV_ROOT", "~/.pex_dev"))
)


@attr.s(frozen=True)
class PythonDistribution(object):
    @classmethod
    def from_venv(cls, venv):
        # type: (str) -> PythonDistribution
        virtualenv = Virtualenv(venv)
        return cls(home=venv, interpreter=virtualenv.interpreter, pip=virtualenv.bin_path("pip"))

    home = attr.ib()  # type: str
    interpreter = attr.ib()  # type: PythonInterpreter
    pip = attr.ib()  # type: str

    @property
    def binary(self):
        # type: () -> str
        return self.interpreter.binary


@attr.s(frozen=True)
class PyenvPythonDistribution(PythonDistribution):
    pyenv_root = attr.ib()  # type: str
    _pyenv_script = attr.ib()  # type: str

    def pyenv_env(self, **extra_env):
        # type: (**str) -> Dict[str, str]
        env = os.environ.copy()
        env.update(extra_env)
        env["PYENV_ROOT"] = self.pyenv_root
        env["PATH"] = os.pathsep.join(
            [os.path.join(self.pyenv_root, path) for path in ("bin", "shims")]
            + os.getenv("PATH", os.defpath).split(os.pathsep)
        )
        return env

    def run_pyenv(
        self,
        args,  # type: Iterable[str]
        **popen_kwargs  # type: Any
    ):
        # type: (...) -> Text
        return cast(
            "Text",
            subprocess.check_output(
                args=[self._pyenv_script] + list(args),
                env=self.pyenv_env(**popen_kwargs.pop("env", {})),
                **popen_kwargs
            ).decode("utf-8"),
        )


def ensure_python_distribution(
    version,  # type: str
    python_version=None,  # type: Optional[str]
    allow_adhoc_version=False,  # type: bool
):
    # type: (...) -> PyenvPythonDistribution
    if not allow_adhoc_version and version not in ALL_PY_VERSIONS:
        raise ValueError("Please constrain version to one of {}".format(ALL_PY_VERSIONS))

    if WINDOWS and _ALL_PY_VERSIONS_TO_VERSION_INFO[version][:2] < (3, 8):
        pytest.skip(
            "Test uses pyenv {version} interpreter which is not supported on Windows.".format(
                version=version
            )
        )

    clone_dir = os.path.abspath(
        os.path.join(PEX_TEST_DEV_ROOT, "pyenv-win" if WINDOWS else "pyenv")
    )
    with atomic_directory(target_dir=clone_dir) as pyenv_root_atomic_dir:
        if not pyenv_root_atomic_dir.is_finalized():
            bootstrap_python_installer(pyenv_root_atomic_dir.work_dir)

    pyenv_root = os.path.join(clone_dir, "pyenv-win") if WINDOWS else clone_dir
    pyenv = os.path.join(pyenv_root, "bin", "pyenv.bat" if WINDOWS else "pyenv")

    interpreter_location = os.path.join(pyenv_root, "versions", version)
    pip = os.path.join(interpreter_location, SCRIPT_DIR, script_name("pip"))

    if WINDOWS:
        python = os.path.join(interpreter_location, "python.exe")
    else:
        major, minor = (python_version or version).split(".")[:2]
        python = os.path.join(
            interpreter_location, "bin", "python{major}.{minor}".format(major=major, minor=minor)
        )

    with atomic_directory(target_dir=interpreter_location) as interpreter_target_dir:
        if not interpreter_target_dir.is_finalized():
            with pyenv_root_atomic_dir.locked():
                subprocess.check_call(args=["git", "pull", "--ff-only"], cwd=pyenv_root)

            env = os.environ.copy()
            env["PYENV_ROOT"] = pyenv_root
            env["CFLAGS"] = "-std=c11"
            if sys.platform.lower().startswith("linux"):
                env["CONFIGURE_OPTS"] = "--enable-shared"
                # The pyenv builder detects `--enable-shared` and sets up `RPATH` via
                # `LDFLAGS=-Wl,-rpath=... $LDFLAGS` to ensure the built python binary links the
                # correct libpython shared lib. Some versions of compiler set the `RUNPATH` instead
                # though which is searched _after_ the `LD_LIBRARY_PATH` environment variable. To
                # ensure an inopportune `LD_LIBRARY_PATH` doesn't fool the pyenv python binary into
                # linking the wrong libpython, force `RPATH`, which is searched 1st by the linker,
                # with `--disable-new-dtags`.
                env["LDFLAGS"] = "-Wl,--disable-new-dtags"
            subprocess.check_call([pyenv, "install", version], env=env)
            subprocess.check_call([python, "-m", "pip", "install", "-U", "pip<22.1"])

    return PyenvPythonDistribution(
        home=interpreter_location,
        interpreter=PythonInterpreter.from_binary(python),
        pip=pip,
        pyenv_root=pyenv_root,
        pyenv_script=pyenv,
    )


def ensure_python_venv(
    version,  # type: str
    latest_pip=True,  # type: bool
    system_site_packages=False,  # type: bool
    tmpdir=None,  # type: Optional[Tempdir]
):
    # type: (...) -> Virtualenv
    pyenv_distribution = ensure_python_distribution(version)
    venv = tmpdir.join("{version}.venv".format(version=version)) if tmpdir else safe_mkdtemp()
    if _ALL_PY_VERSIONS_TO_VERSION_INFO[version][0] == 3:
        args = [pyenv_distribution.binary, "-m", "venv", venv]
        if system_site_packages:
            args.append("--system-site-packages")
        subprocess.check_call(args=args)
    else:
        subprocess.check_call(args=[pyenv_distribution.pip, "install", "virtualenv==16.7.10"])
        args = [pyenv_distribution.binary, "-m", "virtualenv", venv, "-q"]
        if system_site_packages:
            args.append("--system-site-packages")
        subprocess.check_call(args=args)
    python, pip = tuple(os.path.join(venv, "bin", exe) for exe in ("python", "pip"))
    if latest_pip:
        subprocess.check_call(args=[pip, "install", "-U", "pip"])
    return Virtualenv(venv)


def ensure_python_interpreter(version):
    # type: (str) -> str
    return ensure_python_distribution(version).binary


def find_python_interpreter(
    version=(),  # type: Tuple[int, ...]
    implementation=InterpreterImplementation.CPYTHON,  # type: InterpreterImplementation.Value
):
    # type: (...) -> Optional[str]
    for pyenv_version, penv_version_info in _ALL_PY_VERSIONS_TO_VERSION_INFO.items():
        if version and version == penv_version_info[: len(version)]:
            return ensure_python_interpreter(pyenv_version)

    for interpreter in PythonInterpreter.iter():
        if version != interpreter.version[: len(version)]:
            continue
        if implementation is not interpreter.identity.implementation:
            continue
        return interpreter.binary

    return None


def python_venv(
    python,  # type: str
    system_site_packages=False,  # type: bool
    venv_dir=None,  # type: Optional[str]
):
    # type: (...) -> Tuple[str, str]
    venv = Virtualenv.create(
        venv_dir=venv_dir or safe_mkdtemp(),
        interpreter=PythonInterpreter.from_binary(python),
        system_site_packages=system_site_packages,
        install_pip=InstallationChoice.YES,
    )
    return venv.interpreter.binary, venv.bin_path("pip")


def _applicable_py_versions():
    # type: () -> Iterable[str]
    for version in ALL_PY_VERSIONS:
        if WINDOWS and _ALL_PY_VERSIONS_TO_VERSION_INFO[version][:2] < (3, 8):
            continue
        yield version


def all_pythons():
    # type: () -> Tuple[str, ...]
    return tuple(ensure_python_interpreter(version) for version in _applicable_py_versions())


@attr.s(frozen=True)
class VenvFactory(object):
    python_version = attr.ib()  # type: str
    _factory = attr.ib()  # type: Callable[[], Virtualenv]

    def create_venv(self):
        # type: () -> Tuple[str, str]
        venv = self._factory()
        return venv.interpreter.binary, venv.bin_path("pip")


def all_python_venvs(system_site_packages=False):
    # type: (bool) -> Iterable[VenvFactory]
    return tuple(
        VenvFactory(
            python_version=version,
            factory=lambda: ensure_python_venv(version, system_site_packages=system_site_packages),
        )
        for version in _applicable_py_versions()
    )


@contextmanager
def pushd(directory):
    # type: (Text) -> Iterator[None]
    cwd = os.getcwd()
    try:
        os.chdir(directory)
        yield
    finally:
        os.chdir(cwd)


def make_env(**kwargs):
    # type: (**Any) -> Dict[str, str]
    """Create a copy of the current environment with the given modifications.

    The given kwargs add to or update the environment when they have a non-`None` value. When they
    have a `None` value, the environment variable is removed from the environment.

    All non-`None` values are converted to strings by apply `str`.
    """
    env = os.environ.copy()
    env.update((k, str(v)) for k, v in kwargs.items() if v is not None)
    for k, v in kwargs.items():
        if v is None:
            env.pop(k, None)
    return env


def run_commands_with_jitter(
    commands,  # type: Iterable[Iterable[str]]
    path_argument,  # type: str
    extra_env=None,  # type: Optional[Mapping[str, str]]
    delay=2.0,  # type: float
    dest=None,  # type: Optional[str]
):
    # type: (...) -> List[str]
    """Runs the commands with tactics that attempt to introduce randomness in outputs.

    Each command will run against a clean Pex cache with a unique path injected as the value for
    `path_argument`. A unique `PYTHONHASHSEED` is set in the environment for each execution as well.

    Additionally, a delay is inserted between executions. By default, this delay is 2s to ensure zip
    precision is stressed. See: https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT.
    """
    dest_dir = dest or safe_mkdtemp()
    pex_root = os.path.join(dest_dir, "pex_root")

    paths = []
    for index, command in enumerate(commands):
        path = os.path.join(dest_dir, str(index))
        cmd = list(command) + [path_argument, path]

        # Note that we change the `PYTHONHASHSEED` to ensure we're impervious to data structure
        # changes that may result.
        env = make_env(PEX_ROOT=pex_root, PYTHONHASHSEED=(index * 497) + 4)
        if extra_env:
            env.update(extra_env)

        if index > 0:
            safe_sleep(delay)

        # Ensure the PEX is fully rebuilt.
        safe_rmtree(pex_root)
        subprocess.check_call(args=cmd, env=env)
        paths.append(path)
    return paths


def run_command_with_jitter(
    args,  # type: Iterable[str]
    path_argument,  # type: str
    extra_env=None,  # type: Optional[Mapping[str, str]]
    delay=2.0,  # type: float
    count=3,  # type: int
    dest=None,  # type: Optional[str]
):
    # type: (...) -> List[str]
    """Runs the command `count` times in an attempt to introduce randomness.

    Each run of the command will run against a clean Pex cache with a unique path injected as the
    value for `path_argument`. A unique `PYTHONHASHSEED` is set in the environment for each
    execution as well.

    Additionally, a delay is inserted between executions. By default, this delay is 2s to ensure zip
    precision is stressed. See: https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT.
    """
    return run_commands_with_jitter(
        commands=list(itertools.repeat(list(args), count)),
        path_argument=path_argument,
        extra_env=extra_env,
        delay=delay,
        dest=dest,
    )


def pex_project_dir():
    # type: () -> str
    try:
        return os.environ["_PEX_TEST_PROJECT_DIR"]
    except KeyError:
        sys.exit("Pex tests must be run via testing/bin/runtests.py.")


class NonDeterministicWalk:
    """A wrapper around `os.walk` that makes it non-deterministic.

    Makes sure that directories and files are always returned in a different
    order each time it is called.

    Typically used like: `unittest.mock.patch("os.walk", new=NonDeterministicWalk())`
    """

    def __init__(self):
        self._counter = Counter()  # type: Counter[str, int]
        self._original_walk = os.walk

    def __call__(self, *args, **kwargs):
        # type: (*Any, **Any) -> Iterator[Tuple[str, List[str], List[str]]]
        for root, dirs, files in self._original_walk(*args, **kwargs):
            self._increment_counter(root)
            dirs[:] = self._rotate(root, dirs)
            files[:] = self._rotate(root, files)
            yield root, dirs, files

    def _increment_counter(self, counter_key):
        # type: (str) -> int
        self._counter[counter_key] += 1
        return self._counter[counter_key]

    def _rotate(self, counter_key, x):
        # type: (str, List[str]) -> List[str]
        if not x:
            return x
        rotate_by = self._counter[counter_key] % len(x)
        return x[-rotate_by:] + x[:-rotate_by]


def installed_pex_wheel_venv_python(python):
    # type: (str) -> str

    from testing.pex_dist import wheel

    interpreter = PythonInterpreter.from_binary(python)
    pex_wheel = wheel(LocalInterpreter.create(interpreter=interpreter))
    pex_venv_dir = os.path.join(
        PEX_TEST_DEV_ROOT, "pex_venvs", "0", str(interpreter.identity), CacheHelper.hash(pex_wheel)
    )
    with atomic_directory(pex_venv_dir) as atomic_dir:
        if not atomic_dir.is_finalized():
            venv = Virtualenv.create_atomic(atomic_dir, interpreter=interpreter)
            install_wheel_interpreter(pex_wheel, venv.interpreter)
            for _ in venv.rewrite_scripts(
                python=venv.interpreter.binary.replace(atomic_dir.work_dir, atomic_dir.target_dir)
            ):
                # Just ensure the re-writing iterator is driven to completion.
                pass
    return Virtualenv(pex_venv_dir).interpreter.binary
