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

"""pex support for interacting with interpreters."""

from __future__ import absolute_import

import hashlib
import json
import os
import platform
import re
import subprocess
import sys
from textwrap import dedent

from pex import third_party
from pex.common import safe_rmtree
from pex.compatibility import string
from pex.executor import Executor
from pex.jobs import Job, SpawnedJob, execute_parallel
from pex.third_party.packaging import markers, tags
from pex.third_party.pkg_resources import Distribution, Requirement
from pex.tracer import TRACER
from pex.util import CacheHelper
from pex.variables import ENV


class PythonIdentity(object):
  class Error(Exception): pass
  class InvalidError(Error): pass
  class UnknownRequirement(Error): pass

  # TODO(wickman)  Support interpreter-specific versions, e.g. PyPy-2.2.1
  INTERPRETER_NAME_TO_HASHBANG = {
    'CPython': 'python%(major)d.%(minor)d',
    'Jython': 'jython',
    'PyPy': 'pypy',
    'IronPython': 'ipy',
  }

  ABBR_TO_INTERPRETER_NAME = {
    'pp': 'PyPy',
    'jy': 'Jython',
    'ip': 'IronPython',
    'cp': 'CPython',
  }

  @classmethod
  def get(cls):
    supported_tags = tuple(tags.sys_tags())
    preferred_tag = supported_tags[0]
    return cls(
      binary=sys.executable,
      python_tag=preferred_tag.interpreter,
      abi_tag=preferred_tag.abi,
      platform_tag=preferred_tag.platform,
      version=sys.version_info[:3],
      supported_tags=supported_tags,
      env_markers=markers.default_environment()
    )

  @classmethod
  def decode(cls, encoded):
    TRACER.log('creating PythonIdentity from encoded: %s' % encoded, V=9)
    values = json.loads(encoded)
    if len(values) != 7:
      raise cls.InvalidError("Invalid interpreter identity: %s" % encoded)

    supported_tags = values.pop('supported_tags')

    def iter_tags():
      for (interpreter, abi, platform) in supported_tags:
        yield tags.Tag(interpreter=interpreter, abi=abi, platform=platform)

    return cls(supported_tags=iter_tags(), **values)

  @classmethod
  def _find_interpreter_name(cls, python_tag):
    for abbr, interpreter in cls.ABBR_TO_INTERPRETER_NAME.items():
      if python_tag.startswith(abbr):
        return interpreter
    raise ValueError('Unknown interpreter: {}'.format(python_tag))

  def __init__(
      self,
      binary,
      python_tag,
      abi_tag,
      platform_tag,
      version,
      supported_tags,
      env_markers
  ):
    # N.B.: We keep this mapping to support historical values for `distribution` and `requirement`
    # properties.
    self._interpreter_name = self._find_interpreter_name(python_tag)

    self._binary = binary
    self._python_tag = python_tag
    self._abi_tag = abi_tag
    self._platform_tag = platform_tag
    self._version = tuple(version)
    self._supported_tags = tuple(supported_tags)
    self._env_markers = dict(env_markers)

  def encode(self):
    values = dict(
      binary=self._binary,
      python_tag=self._python_tag,
      abi_tag=self._abi_tag,
      platform_tag=self._platform_tag,
      version=self._version,
      supported_tags=[(tag.interpreter, tag.abi, tag.platform) for tag in self._supported_tags],
      env_markers=self._env_markers
    )
    return json.dumps(values)

  @property
  def binary(self):
    return self._binary

  @property
  def python_tag(self):
    return self._python_tag

  @property
  def abi_tag(self):
    return self._abi_tag

  @property
  def platform_tag(self):
    return self._platform_tag

  @property
  def version(self):
    return self._version

  @property
  def version_str(self):
    return '.'.join(map(str, self.version))

  @property
  def supported_tags(self):
    return self._supported_tags

  @property
  def env_markers(self):
    return dict(self._env_markers)

  @property
  def interpreter(self):
    return self._interpreter_name

  @property
  def requirement(self):
    return self.distribution.as_requirement()

  @property
  def distribution(self):
    return Distribution(project_name=self.interpreter, version=self.version_str)

  @classmethod
  def parse_requirement(cls, requirement, default_interpreter='CPython'):
    if isinstance(requirement, Requirement):
      return requirement
    elif isinstance(requirement, string):
      try:
        requirement = Requirement.parse(requirement)
      except ValueError:
        try:
          requirement = Requirement.parse('%s%s' % (default_interpreter, requirement))
        except ValueError:
          raise ValueError('Unknown requirement string: %s' % requirement)
      return requirement
    else:
      raise ValueError('Unknown requirement type: %r' % (requirement,))

  def matches(self, requirement):
    """Given a Requirement, check if this interpreter matches."""
    try:
      requirement = self.parse_requirement(requirement, self._interpreter_name)
    except ValueError as e:
      raise self.UnknownRequirement(str(e))
    return self.distribution in requirement

  def hashbang(self):
    hashbang_string = self.INTERPRETER_NAME_TO_HASHBANG.get(self._interpreter_name, 'CPython') % {
      'major': self._version[0],
      'minor': self._version[1],
      'patch': self._version[2],
    }
    return '#!/usr/bin/env %s' % hashbang_string

  @property
  def python(self):
    # return the python version in the format of the 'python' key for distributions
    # specifically, '2.7', '3.2', etc.
    return '%d.%d' % (self.version[0:2])

  def __str__(self):
    # N.B.: Kept as distinct from __repr__ to support legacy str(identity) used by Pants v1 when
    # forming cache locations.
    return '{interpreter_name}-{major}.{minor}.{patch}'.format(
      interpreter_name=self._interpreter_name,
      major=self._version[0],
      minor=self._version[1],
      patch=self._version[2]
    )

  def __repr__(self):
    return '{type}({binary!r}, {python_tag!r}, {abi_tag!r}, {platform_tag!r}, {version!r})'.format(
      type=self.__class__.__name__,
      binary=self._binary,
      python_tag=self._python_tag,
      abi_tag=self._abi_tag,
      platform_tag=self._platform_tag,
      version=self._version
    )

  def _tup(self):
    return self._binary, self._python_tag, self._abi_tag, self._platform_tag, self._version

  def __eq__(self, other):
    if type(other) is not type(self):
      return NotImplemented
    return self._tup() == other._tup()

  def __hash__(self):
    return hash(self._tup())


class PythonInterpreter(object):
  _REGEXEN = (
    re.compile(r'jython$'),

    # NB: OSX ships python binaries named Python so we allow for capital-P.
    re.compile(r'[Pp]ython$'),

    re.compile(r'python[23]$'),
    re.compile(r'python[23].[0-9]$'),

    # Some distributions include a suffix on the in the interpreter name, similar to PEP-3149
    # E.g. Gentoo has /usr/bin/python3.6m to indicate it was built with pymalloc
    re.compile(r'python[23].[0-9][a-z]$'),

    re.compile(r'pypy$'),
    re.compile(r'pypy-1.[0-9]$'),
  )

  _PYTHON_INTERPRETER_BY_NORMALIZED_PATH = {}

  @staticmethod
  def _normalize_path(path):
    return os.path.realpath(path)

  class Error(Exception): pass
  class IdentificationError(Error): pass
  class InterpreterNotFound(Error): pass

  @classmethod
  def get(cls):
    return cls.from_binary(sys.executable)

  @staticmethod
  def _paths(paths=None):
    return paths or os.getenv('PATH', '').split(os.pathsep)

  @classmethod
  def iter(cls, paths=None):
    """Iterate all interpreters found in `paths`.

    NB: The paths can either be directories to search for python binaries or the paths of python
    binaries themselves.

    :param paths: The paths to look for python interpreters; by default the `PATH`.
    :type paths: list str
    """
    return cls._filter(cls._find(cls._paths(paths=paths)))

  @classmethod
  def all(cls, paths=None):
    return list(cls.iter(paths=paths))

  @classmethod
  def _create_isolated_cmd(cls, binary, args=None, pythonpath=None, env=None):
    cmd = [binary]

    # Don't add the user site directory to `sys.path`.
    #
    # Additionally, it would be nice to pass `-S` to disable adding site-packages but unfortunately
    # some python distributions include portions of the standard library there.
    cmd.append('-s')

    env = cls._sanitized_environment(env=env)
    pythonpath = list(pythonpath or ())
    if pythonpath:
      env['PYTHONPATH'] = os.pathsep.join(pythonpath)
    else:
      # Turn off reading of PYTHON* environment variables.
      cmd.append('-E')

    if args:
      cmd.extend(args)

    rendered_command = ' '.join(cmd)
    if pythonpath:
      rendered_command = 'PYTHONPATH={} {}'.format(env['PYTHONPATH'], rendered_command)
    TRACER.log('Executing: {}'.format(rendered_command), V=3)

    return cmd, env

  @classmethod
  def _execute(cls, binary, args=None, pythonpath=None, env=None, stdin_payload=None, **kwargs):
    cmd, env = cls._create_isolated_cmd(binary, args=args, pythonpath=pythonpath, env=env)
    stdout, stderr = Executor.execute(cmd, stdin_payload=stdin_payload, env=env, **kwargs)
    return cmd, stdout, stderr

  INTERP_INFO_FILE = 'INTERP-INFO'

  @classmethod
  def _spawn_from_binary_external(cls, binary):
    def create_interpreter(stdout):
      identity = stdout.decode('utf-8').strip()
      if not identity:
        raise cls.IdentificationError('Could not establish identity of %s' % binary)
      return cls(PythonIdentity.decode(identity))

    # Part of the PythonInterpreter data are environment markers that depend on the current OS
    # release. That data can change when the OS is upgraded but (some of) the installed interpreters
    # remain the same. As such, include the OS in the hash structure for cached interpreters.
    os_digest = hashlib.sha1()
    for os_identifier in platform.release(), platform.version():
      os_digest.update(os_identifier.encode('utf-8'))
    os_hash = os_digest.hexdigest()

    interpreter_cache_dir = os.path.join(ENV.PEX_ROOT, 'interpreters')
    os_cache_dir = os.path.join(interpreter_cache_dir, os_hash)
    if os.path.isdir(interpreter_cache_dir) and not os.path.isdir(os_cache_dir):
      with TRACER.timed('GCing interpreter cache from prior OS version'):
        safe_rmtree(interpreter_cache_dir)

    interpreter_hash = CacheHelper.hash(binary)
    cache_dir = os.path.join(os_cache_dir, interpreter_hash)
    cache_file = os.path.join(cache_dir, cls.INTERP_INFO_FILE)
    if os.path.isfile(cache_file):
      try:
        with open(cache_file, 'rb') as fp:
          return SpawnedJob.completed(create_interpreter(fp.read()))
      except (IOError, OSError, cls.Error, PythonIdentity.Error):
        safe_rmtree(cache_dir)
        return cls._spawn_from_binary_external(binary)
    else:
      pythonpath = third_party.expose(['pex'])
      cmd, env = cls._create_isolated_cmd(
        binary,
        args=[
          '-c',
          dedent("""\
          import os
          import sys

          from pex.common import atomic_directory, safe_open
          from pex.interpreter import PythonIdentity


          encoded_identity = PythonIdentity.get().encode()
          sys.stdout.write(encoded_identity)
          with atomic_directory({cache_dir!r}) as cache_dir:
            if cache_dir:
              with safe_open(os.path.join(cache_dir, {info_file!r}), 'w') as fp:
                fp.write(encoded_identity)
          """.format(cache_dir=cache_dir, info_file=cls.INTERP_INFO_FILE))
        ],
        pythonpath=pythonpath
      )
      process = Executor.open_process(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
      job = Job(command=cmd, process=process)
      return SpawnedJob.stdout(job, result_func=create_interpreter)

  @classmethod
  def _expand_path(cls, path):
    if os.path.isfile(path):
      return [path]
    elif os.path.isdir(path):
      return sorted(os.path.join(path, fn) for fn in os.listdir(path))
    return []

  @classmethod
  def from_env(cls, hashbang):
    """Resolve a PythonInterpreter as /usr/bin/env would.

    :param hashbang: A string, e.g. "python3.3" representing some binary on the $PATH.
    :return: the first matching interpreter found or `None`.
    :rtype: :class:`PythonInterpreter`
    """
    def hashbang_matches(fn):
      basefile = os.path.basename(fn)
      return hashbang == basefile

    for interpreter in cls._identify_interpreters(filter=hashbang_matches):
      return interpreter

  @classmethod
  def _spawn_from_binary(cls, binary):
    normalized_binary = cls._normalize_path(binary)
    if not os.path.exists(normalized_binary):
      raise cls.InterpreterNotFound(normalized_binary)

    # N.B.: The cache is written as the last step in PythonInterpreter instance initialization.
    cached_interpreter = cls._PYTHON_INTERPRETER_BY_NORMALIZED_PATH.get(normalized_binary)
    if cached_interpreter is not None:
      return SpawnedJob.completed(cached_interpreter)
    if normalized_binary == cls._normalize_path(sys.executable):
      current_interpreter = cls(PythonIdentity.get())
      return SpawnedJob.completed(current_interpreter)
    return cls._spawn_from_binary_external(normalized_binary)

  @classmethod
  def from_binary(cls, binary):
    """Create an interpreter from the given `binary`.

    :param str binary: The path to the python interpreter binary.
    :return: an interpreter created from the given `binary`.
    :rtype: :class:`PythonInterpreter`
    """
    return cls._spawn_from_binary(binary).await_result()

  @classmethod
  def _matches_binary_name(cls, path):
    basefile = os.path.basename(path)
    return any(matcher.match(basefile) is not None for matcher in cls._REGEXEN)

  @classmethod
  def _find(cls, paths):
    """Given a list of files or directories, try to detect python interpreters amongst them.

    Returns an iterator over PythonInterpreter objects.
    """
    return cls._identify_interpreters(filter=cls._matches_binary_name, paths=paths)

  @classmethod
  def _identify_interpreters(cls, filter, paths=None):
    def iter_candidates():
      for path in cls._paths(paths=paths):
        for fn in cls._expand_path(path):
          if filter(fn):
            yield fn

    return execute_parallel(inputs=list(iter_candidates()), spawn_func=cls._spawn_from_binary)

  @classmethod
  def _filter(cls, pythons):
    """Filters duplicate python interpreters and versions we don't support.

    Returns an iterator over PythonInterpreters.
    """
    MAJOR, MINOR, SUBMINOR = range(3)
    def version_filter(version):
      return (version[MAJOR] == 2 and version[MINOR] >= 7 or
              version[MAJOR] == 3 and version[MINOR] >= 5)

    seen = set()
    for interp in pythons:
      version = interp.identity.version
      if version not in seen and version_filter(version):
        seen.add(version)
        yield interp

  @classmethod
  def _sanitized_environment(cls, env=None):
    # N.B. This is merely a hack because sysconfig.py on the default OS X
    # installation of 2.7 breaks.
    env_copy = (env or os.environ).copy()
    env_copy.pop('MACOSX_DEPLOYMENT_TARGET', None)
    return env_copy

  def __init__(self, identity):
    """Construct a PythonInterpreter.

    You should probably use `PythonInterpreter.from_binary` instead.

    :param identity: The :class:`PythonIdentity` of the PythonInterpreter.
    """
    self._identity = identity
    self._binary = self._normalize_path(self.identity.binary)

    self._PYTHON_INTERPRETER_BY_NORMALIZED_PATH[self._binary] = self

  @property
  def binary(self):
    return self._binary

  @property
  def identity(self):
    return self._identity

  @property
  def python(self):
    return self._identity.python

  @property
  def version(self):
    return self._identity.version

  @property
  def version_string(self):
    return str(self._identity)

  def execute(self, args=None, stdin_payload=None, pythonpath=None, env=None, **kwargs):
    return self._execute(self.binary,
                         args=args,
                         stdin_payload=stdin_payload,
                         pythonpath=pythonpath,
                         env=env,
                         **kwargs)

  def open_process(self, args=None, pythonpath=None, env=None, **kwargs):
    cmd, env = self._create_isolated_cmd(self.binary, args=args, pythonpath=pythonpath, env=env)
    process = Executor.open_process(cmd, env=env, **kwargs)
    return cmd, process

  def _tup(self):
    return self._binary, self._identity

  def __hash__(self):
    return hash(self._tup())

  def __eq__(self, other):
    if type(other) is not type(self):
      return NotImplemented
    return self._tup() == other._tup()

  def __lt__(self, other):
    if type(other) is not type(self):
      return NotImplemented
    return self.version < other.version

  def __repr__(self):
    return '{type}({binary!r}, {identity!r})'.format(
      type=self.__class__.__name__,
      binary=self._binary,
      identity=self._identity
    )


def spawn_python_job(
  args,
  env=None,
  interpreter=None,
  expose=None,
  pythonpath=None,
  **subprocess_kwargs
):
  """Spawns a python job.

  :param args: The arguments to pass to the python interpreter.
  :type args: list of str
  :param env: The environment to spawn the python interpreter process in. Defaults to the ambient
              environment.
  :type env: dict of (str, str)
  :param interpreter: The interpreter to use to spawn the python job. Defaults to the current
                      interpreter.
  :type interpreter: :class:`PythonInterpreter`
  :param expose: The names of any vendored distributions to expose to the spawned python process.
                 These will be appended to `pythonpath` if passed.
  :type expose: list of str
  :param pythonpath: The PYTHONPATH to expose to the spawned python process. These will be
                     pre-pended to the `expose` path if passed.
  :type pythonpath: list of str
  :param subprocess_kwargs: Any additional :class:`subprocess.Popen` kwargs to pass through.
  :returns: A job handle to the spawned python process.
  :rtype: :class:`Job`
  """
  pythonpath = list(pythonpath or ())
  if expose:
    subprocess_env = (env or os.environ).copy()
    # In order to expose vendored distributions with their un-vendored import paths in-tact, we
    # need to set `__PEX_UNVENDORED__`. See: vendor.__main__.ImportRewriter._modify_import.
    subprocess_env['__PEX_UNVENDORED__'] = '1'

    pythonpath.extend(third_party.expose(expose))
  else:
    subprocess_env = env

  interpreter = interpreter or PythonInterpreter.get()
  cmd, process = interpreter.open_process(
    args=args,
    pythonpath=pythonpath,
    env=subprocess_env,
    **subprocess_kwargs
  )
  return Job(command=cmd, process=process)
