# Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

import json
import os
import subprocess
from subprocess import CalledProcessError
from textwrap import dedent

import pytest

from pex.common import safe_mkdir, safe_mkdtemp, safe_open, safe_rmtree
from pex.pex_info import PexInfo
from pex.pip import get_pip
from pex.testing import run_pex_command


@pytest.fixture(scope="module")
def pex_project_dir():
    return subprocess.check_output(["git", "rev-parse", "--show-toplevel"]).decode("utf-8").strip()


def test_issues_789_demo(pex_project_dir):
    tmpdir = safe_mkdtemp()

    # 1. Imagine we've pre-resolved the requirements needed in our wheel house.
    requirements = [
        "ansicolors",
        "isort",
        "setuptools",  # N.B.: isort doesn't declare its setuptools dependency.
    ]

    wheelhouse = os.path.join(tmpdir, "wheelhouse")
    get_pip().spawn_download_distributions(
        download_dir=wheelhouse, requirements=requirements
    ).wait()

    # 2. Also imagine this configuration is passed to a tool (PEX or a wrapper as in this test
    # example) via the CLI or other configuration data sources. For example, Pants has a `PythonSetup`
    # that combines with BUILD target data to get you this sort of configuration info outside pex.
    resolver_settings = dict(
        indexes=[],  # Turn off pypi.
        find_links=[wheelhouse],  # Use our wheel house.
        build=False,  # Use only pre-built wheels.
    )

    # 3. That same configuration was used to build a standard pex:
    resolver_args = []
    if len(resolver_settings["find_links"]) == 0:
        resolver_args.append("--no-index")
    else:
        for index in resolver_settings["indexes"]:
            resolver_args.extend(["--index", index])

    for repo in resolver_settings["find_links"]:
        resolver_args.extend(["--find-links", repo])

    resolver_args.append("--build" if resolver_settings["build"] else "--no-build")

    project_code_dir = os.path.join(tmpdir, "project_code_dir")
    with safe_open(os.path.join(project_code_dir, "colorized_isort.py"), "w") as fp:
        fp.write(
            dedent(
                """\
                import colors
                import os
                import subprocess
                import sys


                def run():
                    env = os.environ.copy()
                    env.update(PEX_MODULE='isort')
                    isort_process = subprocess.Popen(
                        sys.argv,
                        env=env,
                        stdout = subprocess.PIPE,
                        stderr = subprocess.PIPE
                    )
                    stdout, stderr = isort_process.communicate()
                    print(colors.green(stdout.decode('utf-8')))
                    print(colors.red(stderr.decode('utf-8')))
                    sys.exit(isort_process.returncode)
    """
            )
        )

    colorized_isort_pex = os.path.join(tmpdir, "colorized_isort.pex")
    args = [
        "--sources-directory",
        project_code_dir,
        "--entry-point",
        "colorized_isort:run",
        "--output-file",
        colorized_isort_pex,
    ]
    result = run_pex_command(args + resolver_args + requirements)
    result.assert_success()

    # 4. Now the tool builds a "dehydrated" PEX using the standard pex + resolve settings as the
    # template.
    ptex_cache = os.path.join(tmpdir, ".ptex")

    colorized_isort_pex_info = PexInfo.from_pex(colorized_isort_pex)
    colorized_isort_pex_info.pex_root = ptex_cache

    # Force the standard pex to extract its code. An external tool like Pants would already know the
    # orignal source code file paths, but we need to discover here.
    colorized_isort_pex_code_dir = os.path.join(
        colorized_isort_pex_info.zip_unsafe_cache, colorized_isort_pex_info.code_hash
    )
    env = os.environ.copy()
    env.update(PEX_ROOT=ptex_cache, PEX_INTERPRETER="1", PEX_FORCE_LOCAL="1")
    subprocess.check_call([colorized_isort_pex, "-c", ""], env=env)

    colorized_isort_ptex_code_dir = os.path.join(tmpdir, "colorized_isort_ptex_code_dir")
    safe_mkdir(colorized_isort_ptex_code_dir)

    code = []
    for root, dirs, files in os.walk(colorized_isort_pex_code_dir):
        rel_root = os.path.relpath(root, colorized_isort_pex_code_dir)
        for f in files:
            # Don't ship compiled python from the code extract above, the target interpreter will not
            # match ours in general.
            if f.endswith(".pyc"):
                continue
            rel_path = os.path.normpath(os.path.join(rel_root, f))
            # The root __main__.py is special for any zipapp including pex, let it write its own
            # __main__.py bootstrap. Similarly. PEX-INFO is special to pex and we want the PEX-INFO for
            # The ptex pex, not the pex being ptexed.
            if rel_path in ("__main__.py", PexInfo.PATH):
                continue
            os.symlink(os.path.join(root, f), os.path.join(colorized_isort_ptex_code_dir, rel_path))
            code.append(rel_path)

    ptex_code_dir = os.path.join(tmpdir, "ptex_code_dir")

    ptex_info = dict(code=code, resolver_settings=resolver_settings)
    with safe_open(os.path.join(ptex_code_dir, "PTEX-INFO"), "w") as fp:
        json.dump(ptex_info, fp)

    with safe_open(os.path.join(ptex_code_dir, "IPEX-INFO"), "w") as fp:
        fp.write(colorized_isort_pex_info.dump())

    with safe_open(os.path.join(ptex_code_dir, "ptex.py"), "w") as fp:
        fp.write(
            dedent(
                """\
                import json
                import os
                import sys

                from pex import resolver
                from pex.common import open_zip
                from pex.pex_builder import PEXBuilder
                from pex.pex_info import PexInfo
                from pex.util import CacheHelper
                from pex.variables import ENV

                self = sys.argv[0]
                ipex_file = '{}.ipex'.format(os.path.splitext(self)[0])

                if not os.path.isfile(ipex_file):
                    print('Hydrating {} to {}'.format(self, ipex_file))

                    ptex_pex_info = PexInfo.from_pex(self)
                    code_root = os.path.join(ptex_pex_info.zip_unsafe_cache, ptex_pex_info.code_hash)
                    with open_zip(self) as zf:
                        # Populate the pex with the pinned requirements and distribution names & hashes.
                        ipex_info = PexInfo.from_json(zf.read('IPEX-INFO'))
                        ipex_builder = PEXBuilder(pex_info=ipex_info)

                        # Populate the pex with the needed code.
                        ptex_info = json.loads(zf.read('PTEX-INFO').decode('utf-8'))
                        for path in ptex_info['code']:
                            ipex_builder.add_source(os.path.join(code_root, path), path)

                    # Perform a fully pinned intransitive resolve to hydrate the install cache (not the
                    # pex!).
                    resolver_settings = ptex_info['resolver_settings']
                    resolved_distributions = resolver.resolve(
                        requirements=[str(req) for req in ipex_info.requirements],
                        cache=ipex_info.pex_root,
                        transitive=False,
                        **resolver_settings
                    )

                    ipex_builder.build(ipex_file)

                os.execv(ipex_file, [ipex_file] + sys.argv[1:])
                """
            )
        )

    colorized_isort_ptex = os.path.join(tmpdir, "colorized_isort.ptex")

    result = run_pex_command(
        [
            "--not-zip-safe",
            "--always-write-cache",
            "--pex-root",
            ptex_cache,
            pex_project_dir,
            "--sources-directory",
            ptex_code_dir,
            "--sources-directory",
            colorized_isort_ptex_code_dir,
            "--entry-point",
            "ptex",
            "--output-file",
            colorized_isort_ptex,
        ]
    )
    result.assert_success()

    subprocess.check_call([colorized_isort_ptex, "--version"])
    with pytest.raises(CalledProcessError):
        subprocess.check_call([colorized_isort_ptex, "--not-a-flag"])

    safe_rmtree(ptex_cache)

    # The dehydrated pex now fails since it lost its hydration from the cache.
    with pytest.raises(CalledProcessError):
        subprocess.check_call([colorized_isort_ptex, "--version"])
