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

from __future__ import absolute_import, print_function

import contextlib
import os
import platform
import random
import subprocess
import sys
from collections import namedtuple
from contextlib import contextmanager
from textwrap import dedent

from pex.common import open_zip, safe_mkdir, safe_mkdtemp, safe_rmtree, temporary_dir, touch
from pex.distribution_target import DistributionTarget
from pex.executor import Executor
from pex.interpreter import PythonInterpreter
from pex.pex import PEX
from pex.pex_builder import PEXBuilder
from pex.pip import get_pip
from pex.third_party.pkg_resources import Distribution
from pex.util import DistributionHelper, named_temporary_file

PY_VER = sys.version_info[:2]
IS_PYPY = hasattr(sys, 'pypy_version_info')
NOT_CPYTHON27 = IS_PYPY or PY_VER != (2, 7)
NOT_CPYTHON36 = IS_PYPY or PY_VER != (3, 6)
IS_LINUX = platform.system() == 'Linux'
IS_NOT_LINUX = not IS_LINUX
NOT_CPYTHON27_OR_OSX = NOT_CPYTHON27 or IS_NOT_LINUX
NOT_CPYTHON36_OR_LINUX = NOT_CPYTHON36 or IS_LINUX


@contextlib.contextmanager
def temporary_filename():
  """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):
  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=''):
  """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):
  """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


def yield_files(directory):
  for root, _, files in os.walk(directory):
    for f in files:
      filename = os.path.join(root, f)
      rel_filename = os.path.relpath(filename, directory)
      yield filename, rel_filename


def write_zipfile(directory, dest, reverse=False):
  with open_zip(dest, 'w') as zf:
    for filename, rel_filename in sorted(yield_files(directory), reverse=reverse):
      zf.write(filename, arcname=rel_filename)
  return dest


@contextlib.contextmanager
def make_project(name='my_project',
                 version='0.0.0',
                 zip_safe=True,
                 install_reqs=None,
                 extras_require=None,
                 entry_points=None):

  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',
          'scripts/shell_script',
        ],
        package_data={%(project_name)r: ['package_data/*.dat']},
        install_requires=%(install_requires)r,
        extras_require=%(extras_require)r,
        entry_points=%(entry_points)r,
      )
    '''),
    'scripts/hello_world': '#!/usr/bin/env python\nprint("hello world!")\n',
    'scripts/shell_script': '#!/usr/bin/env bash\necho hello world\n',
    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,
  }

  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 {}}

  with temporary_content(project_content, interp=interp) as 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, interpreter=None, wheel_dir=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()

  def bdist(self):
    get_pip().spawn_build_wheels(
      distributions=[self._source_dir],
      wheel_dir=self._wheel_dir,
      interpreter=self._interpreter
    ).wait()
    dists = os.listdir(self._wheel_dir)
    if len(dists) == 0:
      raise self.BuildFailure('No distributions were produced!')
    elif len(dists) > 1:
      raise self.BuildFailure('Ambiguous source distributions found: %s' % (' '.join(dists)))
    else:
      return os.path.join(self._wheel_dir, dists[0])


@contextlib.contextmanager
def built_wheel(name='my_project',
                version='0.0.0',
                zip_safe=True,
                install_reqs=None,
                extras_require=None,
                interpreter=None,
                **kwargs):

  with make_project(name=name,
                    version=version,
                    zip_safe=zip_safe,
                    install_reqs=install_reqs,
                    extras_require=extras_require) as td:
    builder = WheelBuilder(td, interpreter=interpreter, **kwargs)
    yield builder.bdist()


@contextlib.contextmanager
def make_source_dir(name='my_project', version='0.0.0', install_reqs=None, extras_require=None):
  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', version='0.0.0', zip_safe=True, interpreter=None, **kwargs):
  with built_wheel(name=name,
                   version=version,
                   zip_safe=zip_safe,
                   interpreter=interpreter,
                   **kwargs) as dist_location:

    install_dir = os.path.join(safe_mkdtemp(), os.path.basename(dist_location))
    get_pip().spawn_install_wheel(
      wheel=dist_location,
      install_dir=install_dir,
      target=DistributionTarget.for_interpreter(interpreter)
    ).wait()
    yield DistributionHelper.distribution_from_path(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,
                     exe_contents=None,
                     dists=None,
                     sources=None,
                     coverage=False,
                     interpreter=None,
                     pex_info=None):
  """Write a pex file that optionally contains an executable entry point.

  :param str td: temporary directory path
  :param str exe_contents: entry point python file
  :param dists: distributions to include, typically sdists or bdists
  :type: list of :class:`pex.third_party.pkg_resources.Distribution`
  :param sources: sources to include, as a list of pairs (env_filename, contents)
  :type sources: list of (str, str)
  :param bool coverage: include coverage header
  :param interpreter: a custom interpreter to use to build the pex
  :type interpreter: :class:`pex.interpreter.PythonInterpreter`
  :param pex_info: a custom PexInfo to use to build the pex.
  :type pex_info: :class:`pex.pex_info.PexInfo`
  """
  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


class IntegResults(namedtuple('results', ['output', 'error', 'return_code'])):
  """Convenience object to return integration run results."""

  def assert_success(self):
    assert self.return_code == 0, ('integration test failed: return_code={}, output={}, error={}'
                                   .format(self.return_code, self.output, self.error))

  def assert_failure(self):
    assert self.return_code != 0


def run_pex_command(args, env=None, python=None, quiet=False):
  """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.
  """
  cmd = [python or sys.executable, '-mpex']
  if not quiet:
    cmd.append('-vvvvv')
  cmd.extend(args)
  process = Executor.open_process(cmd=cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  output, error = process.communicate()
  return IntegResults(output.decode('utf-8'), error.decode('utf-8'), process.returncode)


def run_simple_pex(pex, args=(), interpreter=None, stdin=None, **kwargs):
  p = PEX(pex, interpreter=interpreter)
  process = p.run(args=args,
                  blocking=False,
                  stdin=subprocess.PIPE,
                  stdout=subprocess.PIPE,
                  stderr=subprocess.STDOUT,
                  **kwargs)
  stdout, _ = process.communicate(input=stdin)
  return stdout.replace(b'\r', b''), process.returncode


def run_simple_pex_test(body, args=(), env=None, dists=None, coverage=False, interpreter=None):
  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)


def bootstrap_python_installer(dest):
  safe_rmtree(dest)
  for _ in range(3):
    try:
      subprocess.check_call(
        ['git', 'clone', 'https://github.com/pyenv/pyenv.git', dest]
      )
    except subprocess.CalledProcessError as e:
      print('caught exception: %r' % e)
      continue
    else:
      break
  else:
    raise RuntimeError("Helper method could not clone pyenv from git after 3 tries")
  # Create an empty file indicating the fingerprint of the correct set of test interpreters.
  touch(os.path.join(dest, _INTERPRETER_SET_FINGERPRINT))


# 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.
PY27 = '2.7.15'
PY35 = '3.5.6'
PY36 = '3.6.6'

_VERSIONS = (PY27, PY35, PY36)
# This is the filename of a sentinel file that sits in the pyenv root directory.
# Its purpose is to indicate whether pyenv has the correct interpreters installed
# and will be useful for indicating whether we should trigger a reclone to update
# pyenv.
_INTERPRETER_SET_FINGERPRINT = '_'.join(_VERSIONS) + '_pex_fingerprint'


def ensure_python_distribution(version):
  if version not in _VERSIONS:
    raise ValueError('Please constrain version to one of {}'.format(_VERSIONS))

  pyenv_root = os.path.join(os.getcwd(), '.pyenv_test')
  interpreter_location = os.path.join(pyenv_root, 'versions', version)
  pyenv = os.path.join(pyenv_root, 'bin', 'pyenv')
  pip = os.path.join(interpreter_location, 'bin', 'pip')

  if not os.path.exists(os.path.join(pyenv_root, _INTERPRETER_SET_FINGERPRINT)):
    bootstrap_python_installer(pyenv_root)

  if not os.path.exists(interpreter_location):
    env = os.environ.copy()
    env['PYENV_ROOT'] = pyenv_root
    if sys.platform.lower() == 'linux':
      env['CONFIGURE_OPTS'] = '--enable-shared'
    subprocess.check_call([pyenv, 'install', '--keep', version], env=env)
    subprocess.check_call([pip, 'install', '-U', 'pip'])

  python = os.path.join(interpreter_location, 'bin', 'python' + version[0:3])
  return python, pip


def ensure_python_interpreter(version):
  python, _ = ensure_python_distribution(version)
  return python


@contextmanager
def environment_as(**kwargs):
  existing = {key: os.environ.get(key) for key in kwargs}

  def adjust_environment(mapping):
    for key, value in mapping.items():
      if value is not None:
        os.environ[key] = value
      else:
        del os.environ[key]

  adjust_environment(kwargs)
  try:
    yield
  finally:
    adjust_environment(existing)
