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

import json
import os
import re
import sys
from textwrap import dedent

import pytest

from pex.common import safe_open
from pex.interpreter import PythonInterpreter
from pex.layout import Layout
from pex.os import WINDOWS
from pex.pex_info import PexInfo
from pex.pip.version import PipVersion
from pex.typing import TYPE_CHECKING
from testing import (
    PY310,
    PY311,
    all_pythons,
    ensure_python_interpreter,
    make_env,
    run_pex_command,
    subprocess,
)
from testing.pip import skip_if_only_vendored_pip_supported
from testing.pytest_utils.tmp import Tempdir

if TYPE_CHECKING:
    from typing import Any, Iterable, Iterator, List, Text, Tuple


if WINDOWS:
    pytest.skip(
        "The --sh-boot shebang launch trick only works on systems with shebang support.",
        allow_module_level=True,
    )


@pytest.mark.parametrize(
    ["execution_mode_args"],
    [
        pytest.param([], id="zipapp"),
        pytest.param(["--venv"], id="venv"),
    ],
)
def test_execute(
    tmpdir,  # type: Any
    execution_mode_args,  # type: List[str]
):
    # type: (...) -> None

    cowsay = os.path.join(str(tmpdir), "cowsay.pex")
    run_pex_command(
        args=["cowsay==4.0", "-c", "cowsay", "-o", cowsay, "--sh-boot"] + execution_mode_args
    ).assert_success()
    assert "4.0" == subprocess.check_output(args=[cowsay, "--version"]).decode("utf-8").strip()


@pytest.mark.parametrize(
    ["execution_mode_args"],
    [
        pytest.param([], id="zipapp"),
        pytest.param(["--venv"], id="venv"),
    ],
)
def test_issue_1881(
    tmpdir,  # type: Any
    execution_mode_args,  # type: List[str]
):
    # type: (...) -> None

    pex_root = os.path.join(str(tmpdir), "pex_root")
    os.mkdir(pex_root)
    # Test that the runtime_pex_root is respected even when it is unwritable at build time.
    os.chmod(pex_root, 0o555)
    cowsay = os.path.join(str(tmpdir), "cowsay.pex")
    run_pex_command(
        args=[
            "cowsay==4.0",
            "-c",
            "cowsay",
            "-o",
            cowsay,
            "--python-shebang",
            sys.executable,
            "--sh-boot",
            "--runtime-pex-root",
            pex_root,
        ]
        + execution_mode_args,
        use_pex_whl_venv=sys.version_info[0] == 3,
    ).assert_success()
    # simulate pex_root writable at runtime.
    os.chmod(pex_root, 0o777)

    def _execute_cowsay_pex():
        return (
            subprocess.check_output(
                args=[cowsay, "--version"], env=make_env(PEX_VERBOSE=1), stderr=subprocess.STDOUT
            )
            .decode("utf-8")
            .strip()
        )

    # When this string is logged from the sh_boot script it indicates that the slow
    # path of running the zipapp via python interpreter taken.
    slow_path_output = "Running zipapp pex to lay itself out under PEX_ROOT."
    assert slow_path_output in _execute_cowsay_pex()
    assert slow_path_output not in _execute_cowsay_pex()


def interpreters():
    # type: () -> Iterable[Tuple[Text, List[Text]]]

    def iter_interpreters():
        # type: () -> Iterator[Tuple[Text, List[Text]]]

        def entry(path):
            # type: (Text) -> Tuple[Text, List[Text]]
            return os.path.basename(path), [path]

        yield entry(sys.executable)

        for interpreter in all_pythons():
            yield entry(interpreter)

        locations = (
            subprocess.check_output(
                args=["/usr/bin/env", "bash", "-c", "command -v ash bash busybox dash ksh sh zsh"]
            )
            .decode("utf-8")
            .splitlines()
        )
        for location in locations:
            basename = os.path.basename(location)
            if "busybox" == basename:
                yield "ash (via busybox)", [location, "ash"]
            else:
                yield entry(location)

    return sorted({name: args for name, args in iter_interpreters()}.items())


@pytest.mark.parametrize(
    ["interpreter_cmd"],
    [pytest.param(args, id=name) for name, args in interpreters()],
)
def test_execute_via_interpreter(
    tmpdir,  # type: Any
    interpreter_cmd,  # type: List[str]
):
    # type: (...) -> None

    cowsay = os.path.join(str(tmpdir), "cowsay.pex")
    run_pex_command(
        args=["cowsay==4.0", "-c", "cowsay", "-o", cowsay, "--sh-boot"]
    ).assert_success()

    assert (
        "4.0"
        == subprocess.check_output(args=interpreter_cmd + [cowsay, "--version"])
        .decode("utf-8")
        .strip()
    )


def test_python_shebang_respected(tmpdir):
    # type: (Any) -> None

    cowsay = os.path.join(str(tmpdir), "cowsay.pex")
    run_pex_command(
        args=[
            "cowsay==4.0",
            "-c",
            "cowsay",
            "-o",
            cowsay,
            "--sh-boot",
            "--python-shebang",
            # This is a strange shebang ~no-one would use since it short-circuits the PEX execution
            # to always just print the Python interpreter version, but it serves the purposes of:
            # 1. Proving our python shebang is honored by the bash boot.
            # 2. The bash boot treatment can handle shebangs with arguments in them.
            "{python} -V".format(python=sys.executable),
        ]
    ).assert_success()

    # N.B.: Python 2.7 does not send version to stdout; so we redirect stdout to stderr to be able
    # to uniformly retrieve the Python version.
    output = subprocess.check_output(args=[cowsay], stderr=subprocess.STDOUT).decode("utf-8")
    version = "Python {version}".format(version=".".join(map(str, sys.version_info[:3])))
    assert output.startswith(version), output


execution_mode = pytest.mark.parametrize(
    "execution_mode_args",
    [
        pytest.param([], id="ZIPAPP"),
        pytest.param(["--venv"], id="VENV"),
        pytest.param(["--sh-boot"], id="ZIPAPP (--sh-boot)"),
        pytest.param(["--venv", "--sh-boot"], id="VENV (--sh-boot)"),
    ],
)
layouts = pytest.mark.parametrize(
    "layout", [pytest.param(layout, id=layout.value) for layout in Layout.values()]
)


@execution_mode
@layouts
@skip_if_only_vendored_pip_supported
def test_issue_1782(
    tmpdir,  # type: Tempdir
    pex_wheel,  # type: str
    execution_mode_args,  # type: List[str]
    layout,  # type: Layout.Value
):
    # type: (...) -> None

    pex = os.path.realpath(tmpdir.join("pex.sh"))
    pex_exe = pex if layout is Layout.ZIPAPP else os.path.join(pex, "pex")

    pex_root = os.path.realpath(tmpdir.join("pex_root"))
    python = "python{major}.{minor}".format(major=sys.version_info[0], minor=sys.version_info[1])
    run_pex_command(
        args=[
            "--pex-root",
            pex_root,
            "--runtime-pex-root",
            pex_root,
            "--pip-version",
            PipVersion.LATEST_COMPATIBLE.value,
            pex_wheel,
            "-c",
            "pex",
            "-o",
            pex,
            "--python-shebang",
            "/usr/bin/env {python}".format(python=python),
            "--layout",
            str(layout),
        ]
        + execution_mode_args
    ).assert_success()

    usage_line_re = re.compile(r"^usage: {argv0}".format(argv0=re.escape(pex_exe)))
    help_line1 = (
        subprocess.check_output(
            args=[pex_exe, "-h"], env=make_env(COLUMNS=max(80, len(usage_line_re.pattern) + 10))
        )
        .decode("utf-8")
        .splitlines()[0]
        .strip()
    )
    assert usage_line_re.match(help_line1), (
        "\n"
        "expected: {expected}\n"
        "actual:   {actual}".format(expected=usage_line_re.pattern, actual=help_line1)
    )
    assert (
        pex
        == subprocess.check_output(
            args=[pex_exe, "-c", "import os; print(os.environ['PEX'])"],
            env=make_env(PEX_INTERPRETER=1),
        )
        .decode("utf-8")
        .strip()
    )


@execution_mode
@layouts
def test_argv0(
    tmpdir,  # type: Any
    execution_mode_args,  # type: List[str]
    layout,  # type: Layout.Value
):
    # type: (...) -> None

    pex = os.path.realpath(os.path.join(str(tmpdir), "pex.sh"))
    pex_exe = pex if layout is Layout.ZIPAPP else os.path.join(pex, "pex")

    src = os.path.join(str(tmpdir), "src")
    with safe_open(os.path.join(src, "app.py"), "w") as fp:
        fp.write(
            dedent(
                """\
                import json
                import os
                import sys


                def main():
                    print(json.dumps({"PEX": os.environ["PEX"], "argv0": sys.argv[0]}))


                if __name__ == "__main__":
                    main()
                """
            )
        )

    run_pex_command(
        args=["-D", src, "-e", "app:main", "-o", pex, "--layout", str(layout)] + execution_mode_args
    ).assert_success()
    assert {"PEX": pex, "argv0": pex_exe} == json.loads(subprocess.check_output(args=[pex_exe]))

    run_pex_command(
        args=["-D", src, "-m", "app", "-o", pex, "--layout", str(layout)] + execution_mode_args
    ).assert_success()
    data = json.loads(subprocess.check_output(args=[pex_exe]))
    assert pex == data.pop("PEX")
    assert "app.py" == os.path.basename(data.pop("argv0")), (
        "When executing modules we expect runpy.run_module to `alter_sys` in order to support "
        "pickling and other use cases as outlined in https://github.com/pex-tool/pex/issues/1018."
    )
    assert {} == data


def create_empty_pex(
    tmpdir,  # type: Tempdir
    extra_args,  # type: List[str]
):
    # type: (...) -> str
    pex_root = tmpdir.join("pex-root")
    pex = os.path.realpath(tmpdir.join("empty.pex"))
    run_pex_command(
        args=(["--pex-root", pex_root, "--runtime-pex-root", pex_root, "-o", pex] + extra_args)
    ).assert_success()
    return pex


@execution_mode
def test_issue_2726_pex_tools(
    tmpdir,  # type: Tempdir
    execution_mode_args,  # type: List[str]
):
    # type: (...) -> None

    pex = create_empty_pex(tmpdir, execution_mode_args + ["--include-tools"])

    def _check_pex_tools():
        output = subprocess.check_output(args=[pex, "info"], env=make_env(PEX_TOOLS=1))
        # Check for output that implies we got sensible `PEX_TOOLS=1 ./some.pex info` output:
        assert PexInfo.from_pex(pex).as_json_dict() == json.loads(output)

    # run the tools with an empty/unseeded PEX_ROOT
    _check_pex_tools()

    # ensure the PEX_ROOT is fully seeded e.g. with a venv for venv execution mode
    subprocess.check_output(args=[pex, "-c", ""])

    # run the tools with seeded PEX_ROOT
    _check_pex_tools()


@execution_mode
@pytest.mark.parametrize("pex_env_var", ["PEX_PYTHON", "PEX_PYTHON_PATH"])
def test_issue_2728_pex_python_and_pex_python_path(
    tmpdir,  # type: Tempdir
    execution_mode_args,  # type: List[str]
    pex_env_var,
):
    # type: (...) -> None

    pex = create_empty_pex(tmpdir, execution_mode_args)

    other_python = ensure_python_interpreter(PY310 if sys.version_info[:2] == (3, 11) else PY311)

    def _check_pex_with_python():
        output = subprocess.check_output(
            args=[pex, "-c", "import sys; print('.'.join(map(str, sys.version_info[:3])))"],
            env=make_env(**{pex_env_var: other_python}),
        )
        assert (
            ".".join(map(str, PythonInterpreter.from_binary(other_python).version[:3]))
            == output.decode("utf-8").strip()
        )

    _check_pex_with_python()

    # ensure the PEX_ROOT is fully seeded e.g. with a venv for venv execution mode
    subprocess.check_output(args=[pex, "-c", ""])

    _check_pex_with_python()


@execution_mode
def test_issue_2728_pex_path(
    tmpdir,  # type: Tempdir
    execution_mode_args,  # type: List[str]
):
    # type: (...) -> None

    pex = create_empty_pex(tmpdir, execution_mode_args)

    cowsay = tmpdir.join("cowsay.pex")
    run_pex_command(args=["cowsay==5.0", "-o", cowsay]).assert_success()

    def _check_pex_with_path():
        output = subprocess.check_output(
            args=[pex, "-c", "import cowsay; print(cowsay.__version__)"],
            env=make_env(PEX_PATH=cowsay),
        )
        assert b"5.0\n" == output

    _check_pex_with_path()

    # ensure the PEX_ROOT is fully seeded e.g. with a venv for venv execution mode
    subprocess.check_output(args=[pex, "-c", ""])

    _check_pex_with_path()
