# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Support for installing custom code and required dependencies.

Workflows, with the exception of very simple ones, are organized in multiple
modules and packages. Typically, these modules and packages have
dependencies on other standard libraries. Beam relies on the Python
setuptools package to handle these scenarios. For further details please read:
https://pythonhosted.org/an_example_pypi_project/setuptools.html

When a runner tries to run a pipeline it will check for a --requirements_file
and a --setup_file option.

If --setup_file is present then it is assumed that the folder containing the
file specified by the option has the typical layout required by setuptools and
it will run 'python setup.py sdist' to produce a source distribution. The
resulting tarball (a .tar or .tar.gz file) will be staged at the staging
location specified as job option. When a worker starts it will check for the
presence of this file and will run 'easy_install tarball' to install the
package in the worker.

If --requirements_file is present then the file specified by the option will be
staged in the staging location.  When a worker starts it will check for the
presence of this file and will run 'pip install -r requirements.txt'. A
requirements file can be easily generated by running 'pip freeze -r
requirements.txt'. The reason a runner does not run this automatically is
because quite often only a small fraction of the dependencies present in a
requirements.txt file are actually needed for remote execution and therefore a
one-time manual trimming is desirable.

TODO(silviuc): Should we allow several setup packages?
TODO(silviuc): We should allow customizing the exact command for setup build.
"""
from __future__ import absolute_import

import glob
import logging
import os
import shutil
import sys
import tempfile

import pkg_resources

from apache_beam.internal import pickler
from apache_beam.internal.http_client import get_new_http
from apache_beam.io.filesystems import FileSystems
from apache_beam.options.pipeline_options import DebugOptions
from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.options.pipeline_options import WorkerOptions
# TODO(angoenka): Remove reference to dataflow internal names
from apache_beam.runners.dataflow.internal.names import DATAFLOW_SDK_TARBALL_FILE
from apache_beam.runners.internal import names
from apache_beam.utils import processes
from apache_beam.utils import retry

# All constants are for internal use only; no backwards-compatibility
# guarantees.

# Standard file names used for staging files.
WORKFLOW_TARBALL_FILE = 'workflow.tar.gz'
REQUIREMENTS_FILE = 'requirements.txt'
EXTRA_PACKAGES_FILE = 'extra_packages.txt'


def retry_on_non_zero_exit(exception):
  if (isinstance(exception, processes.CalledProcessError) and
      exception.returncode != 0):
    return True
  return False


class Stager(object):
  """Abstract Stager identifies and copies the appropriate artifacts to the
  staging location.
  Implementation of this stager has to implement :func:`stage_artifact` and
  :func:`commit_manifest`.
  """

  def stage_artifact(self, local_path_to_artifact, artifact_name):
    """ Stages the artifact to Stager._staging_location and adds artifact_name
        to the manifest of artifacts that have been staged."""
    raise NotImplementedError

  def commit_manifest(self):
    """Commits manifest."""
    raise NotImplementedError

  @staticmethod
  def get_sdk_package_name():
    """For internal use only; no backwards-compatibility guarantees.
        Returns the PyPI package name to be staged."""
    return names.BEAM_PACKAGE_NAME

  def stage_job_resources(self,
                          options,
                          build_setup_args=None,
                          temp_dir=None,
                          populate_requirements_cache=None,
                          staging_location=None):
    """For internal use only; no backwards-compatibility guarantees.

        Creates (if needed) and stages job resources to staging_location.

        Args:
          options: Command line options. More specifically the function will
            expect requirements_file, setup_file, and save_main_session options
            to be present.
          build_setup_args: A list of command line arguments used to build a
            setup package. Used only if options.setup_file is not None. Used
            only for testing.
          temp_dir: Temporary folder where the resource building can happen. If
            None then a unique temp directory will be created. Used only for
            testing.
          populate_requirements_cache: Callable for populating the requirements
            cache. Used only for testing.
          staging_location: Location to stage the file.

        Returns:
          A list of file names (no paths) for the resources staged. All the
          files are assumed to be staged at staging_location.

        Raises:
          RuntimeError: If files specified are not found or error encountered
          while trying to create the resources (e.g., build a setup package).
        """
    temp_dir = temp_dir or tempfile.mkdtemp()
    resources = []

    setup_options = options.view_as(SetupOptions)
    # Make sure that all required options are specified.
    if staging_location is None:
      raise RuntimeError('The staging_location must be specified.')

    # Stage a requirements file if present.
    if setup_options.requirements_file is not None:
      if not os.path.isfile(setup_options.requirements_file):
        raise RuntimeError(
            'The file %s cannot be found. It was specified in the '
            '--requirements_file command line option.' %
            setup_options.requirements_file)
      staged_path = FileSystems.join(staging_location, REQUIREMENTS_FILE)
      self.stage_artifact(setup_options.requirements_file, staged_path)
      resources.append(REQUIREMENTS_FILE)
      requirements_cache_path = (
          os.path.join(tempfile.gettempdir(), 'dataflow-requirements-cache')
          if setup_options.requirements_cache is None else
          setup_options.requirements_cache)
      # Populate cache with packages from requirements and stage the files
      # in the cache.
      if not os.path.exists(requirements_cache_path):
        os.makedirs(requirements_cache_path)
      (populate_requirements_cache if populate_requirements_cache else
       Stager._populate_requirements_cache)(setup_options.requirements_file,
                                            requirements_cache_path)
      for pkg in glob.glob(os.path.join(requirements_cache_path, '*')):
        self.stage_artifact(
            pkg, FileSystems.join(staging_location, os.path.basename(pkg)))
        resources.append(os.path.basename(pkg))

    # Handle a setup file if present.
    # We will build the setup package locally and then copy it to the staging
    # location because the staging location is a remote path and the file cannot
    # be created directly there.
    if setup_options.setup_file is not None:
      if not os.path.isfile(setup_options.setup_file):
        raise RuntimeError(
            'The file %s cannot be found. It was specified in the '
            '--setup_file command line option.' % setup_options.setup_file)
      if os.path.basename(setup_options.setup_file) != 'setup.py':
        raise RuntimeError(
            'The --setup_file option expects the full path to a file named '
            'setup.py instead of %s' % setup_options.setup_file)
      tarball_file = Stager._build_setup_package(setup_options.setup_file,
                                                 temp_dir, build_setup_args)
      staged_path = FileSystems.join(staging_location, WORKFLOW_TARBALL_FILE)
      self.stage_artifact(tarball_file, staged_path)
      resources.append(WORKFLOW_TARBALL_FILE)

    # Handle extra local packages that should be staged.
    if setup_options.extra_packages is not None:
      resources.extend(
          self._stage_extra_packages(
              setup_options.extra_packages, staging_location,
              temp_dir=temp_dir))

    # Handle jar packages that should be staged for Java SDK Harness.
    jar_packages = options.view_as(
        DebugOptions).lookup_experiment('jar_packages')
    if jar_packages is not None:
      resources.extend(
          self._stage_jar_packages(
              jar_packages.split(','), staging_location,
              temp_dir=temp_dir))

    # Pickle the main session if requested.
    # We will create the pickled main session locally and then copy it to the
    # staging location because the staging location is a remote path and the
    # file cannot be created directly there.
    if setup_options.save_main_session:
      pickled_session_file = os.path.join(temp_dir,
                                          names.PICKLED_MAIN_SESSION_FILE)
      pickler.dump_session(pickled_session_file)
      staged_path = FileSystems.join(staging_location,
                                     names.PICKLED_MAIN_SESSION_FILE)
      self.stage_artifact(pickled_session_file, staged_path)
      resources.append(names.PICKLED_MAIN_SESSION_FILE)

    if hasattr(setup_options, 'sdk_location'):

      if (setup_options.sdk_location == 'default') or Stager._is_remote_path(
          setup_options.sdk_location):
        # If --sdk_location is not specified then the appropriate package
        # will be obtained from PyPI (https://pypi.python.org) based on the
        # version of the currently running SDK. If the option is
        # present then no version matching is made and the exact URL or path
        # is expected.
        #
        # Unit tests running in the 'python setup.py test' context will
        # not have the sdk_location attribute present and therefore we
        # will not stage SDK.
        sdk_remote_location = 'pypi' if (setup_options.sdk_location == 'default'
                                        ) else setup_options.sdk_location
        resources.extend(
            self._stage_beam_sdk(sdk_remote_location, staging_location,
                                 temp_dir))
      elif setup_options.sdk_location == 'container':
        # Use the SDK that's built into the container, rather than re-staging
        # it.
        pass
      else:
        # This branch is also used by internal tests running with the SDK built
        # at head.
        if os.path.isdir(setup_options.sdk_location):
          # TODO(angoenka): remove reference to Dataflow
          sdk_path = os.path.join(setup_options.sdk_location,
                                  DATAFLOW_SDK_TARBALL_FILE)
        else:
          sdk_path = setup_options.sdk_location

        if os.path.isfile(sdk_path):
          logging.info('Copying Beam SDK "%s" to staging location.', sdk_path)
          staged_path = FileSystems.join(
              staging_location,
              Stager._desired_sdk_filename_in_staging_location(
                  setup_options.sdk_location))
          self.stage_artifact(sdk_path, staged_path)
          _, sdk_staged_filename = FileSystems.split(staged_path)
          resources.append(sdk_staged_filename)
        else:
          if setup_options.sdk_location == 'default':
            raise RuntimeError('Cannot find default Beam SDK tar file "%s"'
                               % sdk_path)
          elif not setup_options.sdk_location:
            logging.info('Beam SDK will not be staged since --sdk_location '
                         'is empty.')
          else:
            raise RuntimeError(
                'The file "%s" cannot be found. Its location was specified by '
                'the --sdk_location command-line option.' % sdk_path)

    worker_options = options.view_as(WorkerOptions)
    dataflow_worker_jar = getattr(worker_options, 'dataflow_worker_jar', None)
    if dataflow_worker_jar is not None:
      jar_staged_filename = 'dataflow-worker.jar'
      staged_path = FileSystems.join(staging_location, jar_staged_filename)
      self.stage_artifact(dataflow_worker_jar, staged_path)
      resources.append(jar_staged_filename)

    # Delete all temp files created while staging job resources.
    shutil.rmtree(temp_dir)
    retrieval_token = self.commit_manifest()
    return retrieval_token, resources

  @staticmethod
  def _download_file(from_url, to_path):
    """Downloads a file over http/https from a url or copy it from a remote
        path to local path."""
    if from_url.startswith('http://') or from_url.startswith('https://'):
      # TODO(silviuc): We should cache downloads so we do not do it for every
      # job.
      try:
        # We check if the file is actually there because wget returns a file
        # even for a 404 response (file will contain the contents of the 404
        # response).
        # TODO(angoenka): Extract and use the filename when downloading file.
        response, content = get_new_http().request(from_url)
        if int(response['status']) >= 400:
          raise RuntimeError(
              'Artifact not found at %s (response: %s)' % (from_url, response))
        with open(to_path, 'w') as f:
          f.write(content)
      except Exception:
        logging.info('Failed to download Artifact from %s', from_url)
        raise
    else:
      if not os.path.isdir(os.path.dirname(to_path)):
        logging.info(
            'Created folder (since we have not done yet, and any errors '
            'will follow): %s ', os.path.dirname(to_path))
        os.mkdir(os.path.dirname(to_path))
      shutil.copyfile(from_url, to_path)

  @staticmethod
  def _is_remote_path(path):
    return path.find('://') != -1

  def _stage_jar_packages(self, jar_packages, staging_location, temp_dir):
    """Stages a list of local jar packages for Java SDK Harness.

    :param jar_packages: Ordered list of local paths to jar packages to be
      staged. Only packages on localfile system and GCS are supported.
    :param staging_location: Staging location for the packages.
    :param temp_dir: Temporary folder where the resource building can happen.
    :return: A list of file names (no paths) for the resource staged. All the
      files are assumed to be staged in staging_location.
    :raises:
      RuntimeError: If files specified are not found or do not have expected
        name patterns.
    """
    resources = []
    staging_temp_dir = tempfile.mkdtemp(dir=temp_dir)
    local_packages = []
    for package in jar_packages:
      if not os.path.basename(package).endswith('.jar'):
        raise RuntimeError(
            'The --experiment=\'jar_packages=\' option expects a full path '
            'ending with ".jar" instead of %s' % package)

      if not os.path.isfile(package):
        if Stager._is_remote_path(package):
          # Download remote package.
          logging.info('Downloading jar package: %s locally before staging',
                       package)
          _, last_component = FileSystems.split(package)
          local_file_path = FileSystems.join(staging_temp_dir, last_component)
          Stager._download_file(package, local_file_path)
        else:
          raise RuntimeError(
              'The file %s cannot be found. It was specified in the '
              '--experiment=\'jar_packages=\' command line option.' % package)
      else:
        local_packages.append(package)

    local_packages.extend([
        FileSystems.join(staging_temp_dir, f)
        for f in os.listdir(staging_temp_dir)
    ])

    for package in local_packages:
      basename = os.path.basename(package)
      staged_path = FileSystems.join(staging_location, basename)
      self.stage_artifact(package, staged_path)
      resources.append(basename)

    return resources

  def _stage_extra_packages(self, extra_packages, staging_location, temp_dir):
    """Stages a list of local extra packages.

      Args:
        extra_packages: Ordered list of local paths to extra packages to be
          staged. Only packages on localfile system and GCS are supported.
        staging_location: Staging location for the packages.
        temp_dir: Temporary folder where the resource building can happen.
          Caller is responsible for cleaning up this folder after this function
          returns.

      Returns:
        A list of file names (no paths) for the resources staged. All the files
        are assumed to be staged in staging_location.

      Raises:
        RuntimeError: If files specified are not found or do not have expected
          name patterns.
      """
    resources = []
    staging_temp_dir = tempfile.mkdtemp(dir=temp_dir)
    local_packages = []
    for package in extra_packages:
      if not (os.path.basename(package).endswith('.tar') or
              os.path.basename(package).endswith('.tar.gz') or
              os.path.basename(package).endswith('.whl') or
              os.path.basename(package).endswith('.zip')):
        raise RuntimeError(
            'The --extra_package option expects a full path ending with '
            '".tar", ".tar.gz", ".whl" or ".zip" instead of %s' % package)
      if os.path.basename(package).endswith('.whl'):
        logging.warning(
            'The .whl package "%s" is provided in --extra_package. '
            'This functionality is not officially supported. Since wheel '
            'packages are binary distributions, this package must be '
            'binary-compatible with the worker environment (e.g. Python 2.7 '
            'running on an x64 Linux host).')

      if not os.path.isfile(package):
        if Stager._is_remote_path(package):
          # Download remote package.
          logging.info('Downloading extra package: %s locally before staging',
                       package)
          _, last_component = FileSystems.split(package)
          local_file_path = FileSystems.join(staging_temp_dir, last_component)
          Stager._download_file(package, local_file_path)
        else:
          raise RuntimeError(
              'The file %s cannot be found. It was specified in the '
              '--extra_packages command line option.' % package)
      else:
        local_packages.append(package)

    local_packages.extend([
        FileSystems.join(staging_temp_dir, f)
        for f in os.listdir(staging_temp_dir)
    ])

    for package in local_packages:
      basename = os.path.basename(package)
      staged_path = FileSystems.join(staging_location, basename)
      self.stage_artifact(package, staged_path)
      resources.append(basename)
    # Create a file containing the list of extra packages and stage it.
    # The file is important so that in the worker the packages are installed
    # exactly in the order specified. This approach will avoid extra PyPI
    # requests. For example if package A depends on package B and package A
    # is installed first then the installer will try to satisfy the
    # dependency on B by downloading the package from PyPI. If package B is
    # installed first this is avoided.
    with open(os.path.join(temp_dir, EXTRA_PACKAGES_FILE), 'wt') as f:
      for package in local_packages:
        f.write('%s\n' % os.path.basename(package))
    staged_path = FileSystems.join(staging_location, EXTRA_PACKAGES_FILE)
    # Note that the caller of this function is responsible for deleting the
    # temporary folder where all temp files are created, including this one.
    self.stage_artifact(
        os.path.join(temp_dir, EXTRA_PACKAGES_FILE), staged_path)
    resources.append(EXTRA_PACKAGES_FILE)

    return resources

  @staticmethod
  def _get_python_executable():
    # Allow overriding the python executable to use for downloading and
    # installing dependencies, otherwise use the python executable for
    # the current process.
    python_bin = os.environ.get('BEAM_PYTHON') or sys.executable
    if not python_bin:
      raise ValueError('Could not find Python executable.')
    return python_bin

  @staticmethod
  @retry.with_exponential_backoff(num_retries=4,
                                  retry_filter=retry_on_non_zero_exit)
  def _populate_requirements_cache(requirements_file, cache_dir):
    # The 'pip download' command will not download again if it finds the
    # tarball with the proper version already present.
    # It will get the packages downloaded in the order they are presented in
    # the requirements file and will not download package dependencies.
    cmd_args = [
        Stager._get_python_executable(),
        '-m',
        'pip',
        'download',
        '--dest',
        cache_dir,
        '-r',
        requirements_file,
        '--exists-action',
        'i',
        # Download from PyPI source distributions.
        '--no-binary',
        ':all:'
    ]
    logging.info('Executing command: %s', cmd_args)
    processes.check_output(cmd_args, stderr=processes.STDOUT)

  @staticmethod
  def _build_setup_package(setup_file, temp_dir, build_setup_args=None):
    saved_current_directory = os.getcwd()
    try:
      os.chdir(os.path.dirname(setup_file))
      if build_setup_args is None:
        build_setup_args = [
            Stager._get_python_executable(),
            os.path.basename(setup_file), 'sdist', '--dist-dir', temp_dir
        ]
      logging.info('Executing command: %s', build_setup_args)
      processes.check_output(build_setup_args)
      output_files = glob.glob(os.path.join(temp_dir, '*.tar.gz'))
      if not output_files:
        raise RuntimeError(
            'File %s not found.' % os.path.join(temp_dir, '*.tar.gz'))
      return output_files[0]
    finally:
      os.chdir(saved_current_directory)

  @staticmethod
  def _desired_sdk_filename_in_staging_location(sdk_location):
    """Returns the name that SDK file should have in the staging location.
      Args:
        sdk_location: Full path to SDK file.
      """
    if sdk_location.endswith('.whl'):
      _, wheel_filename = FileSystems.split(sdk_location)
      if wheel_filename.startswith('apache_beam'):
        return wheel_filename
      else:
        raise RuntimeError('Unrecognized SDK wheel file: %s' % sdk_location)
    else:
      return DATAFLOW_SDK_TARBALL_FILE

  def _stage_beam_sdk(self, sdk_remote_location, staging_location, temp_dir):
    """Stages a Beam SDK file with the appropriate version.

      Args:
        sdk_remote_location: A URL from which thefile can be downloaded or a
          remote file location. The SDK file can be a tarball or a wheel. Set
          to 'pypi' to download and stage a wheel and source SDK from PyPi.
        staging_location: Location where the SDK file should be copied.
        temp_dir: path to temporary location where the file should be
          downloaded.

      Returns:
        A list of SDK files that were staged to the staging location.

      Raises:
        RuntimeError: if staging was not successful.
      """
    if sdk_remote_location == 'pypi':
      sdk_local_file = Stager._download_pypi_sdk_package(temp_dir)
      sdk_sources_staged_name = Stager.\
          _desired_sdk_filename_in_staging_location(sdk_local_file)
      staged_path = FileSystems.join(staging_location, sdk_sources_staged_name)
      logging.info('Staging SDK sources from PyPI to %s', staged_path)
      self.stage_artifact(sdk_local_file, staged_path)
      staged_sdk_files = [sdk_sources_staged_name]
      try:
        # Stage binary distribution of the SDK, for now on a best-effort basis.
        sdk_local_file = Stager._download_pypi_sdk_package(
            temp_dir, fetch_binary=True,
            language_version_tag='%d%d' % (sys.version_info[0],
                                           sys.version_info[1]),
            abi_tag='cp%d%d%s' % (
                sys.version_info[0],
                sys.version_info[1],
                'mu' if sys.version_info[0] < 3 else 'm'))
        sdk_binary_staged_name = Stager.\
            _desired_sdk_filename_in_staging_location(sdk_local_file)
        staged_path = FileSystems.join(staging_location, sdk_binary_staged_name)
        logging.info('Staging binary distribution of the SDK from PyPI to %s',
                     staged_path)
        self.stage_artifact(sdk_local_file, staged_path)
        staged_sdk_files.append(sdk_binary_staged_name)
      except RuntimeError as e:
        logging.warn(
            'Failed to download requested binary distribution '
            'of the SDK: %s', repr(e))

      return staged_sdk_files
    elif Stager._is_remote_path(sdk_remote_location):
      local_download_file = os.path.join(temp_dir, 'beam-sdk.tar.gz')
      Stager._download_file(sdk_remote_location, local_download_file)
      staged_name = Stager._desired_sdk_filename_in_staging_location(
          sdk_remote_location)
      staged_path = FileSystems.join(staging_location, staged_name)
      logging.info('Staging Beam SDK from %s to %s', sdk_remote_location,
                   staged_path)
      self.stage_artifact(local_download_file, staged_path)
      return [staged_name]
    else:
      raise RuntimeError(
          'The --sdk_location option was used with an unsupported '
          'type of location: %s' % sdk_remote_location)

  @staticmethod
  def _download_pypi_sdk_package(temp_dir,
                                 fetch_binary=False,
                                 language_version_tag='27',
                                 language_implementation_tag='cp',
                                 abi_tag='cp27mu',
                                 platform_tag='manylinux1_x86_64'):
    """Downloads SDK package from PyPI and returns path to local path."""
    package_name = Stager.get_sdk_package_name()
    try:
      version = pkg_resources.get_distribution(package_name).version
    except pkg_resources.DistributionNotFound:
      raise RuntimeError('Please set --sdk_location command-line option '
                         'or install a valid {} distribution.'
                         .format(package_name))
    cmd_args = [
        Stager._get_python_executable(), '-m', 'pip', 'download', '--dest',
        temp_dir,
        '%s==%s' % (package_name, version), '--no-deps'
    ]

    if fetch_binary:
      logging.info('Downloading binary distribution of the SDK from PyPi')
      # Get a wheel distribution for the SDK from PyPI.
      cmd_args.extend([
          '--only-binary', ':all:', '--python-version', language_version_tag,
          '--implementation', language_implementation_tag, '--abi', abi_tag,
          '--platform', platform_tag
      ])
      # Example wheel: apache_beam-2.4.0-cp27-cp27mu-manylinux1_x86_64.whl
      expected_files = [
          os.path.join(
              temp_dir, '%s-%s-%s%s-%s-%s.whl' % (package_name.replace(
                  '-', '_'), version, language_implementation_tag,
                                                  language_version_tag, abi_tag,
                                                  platform_tag))
      ]
    else:
      logging.info('Downloading source distribution of the SDK from PyPi')
      cmd_args.extend(['--no-binary', ':all:'])
      expected_files = [
          os.path.join(temp_dir, '%s-%s.zip' % (package_name, version)),
          os.path.join(temp_dir, '%s-%s.tar.gz' % (package_name, version))
      ]

    logging.info('Executing command: %s', cmd_args)
    try:
      processes.check_output(cmd_args)
    except processes.CalledProcessError as e:
      raise RuntimeError(repr(e))

    for sdk_file in expected_files:
      if os.path.exists(sdk_file):
        return sdk_file

    raise RuntimeError(
        'Failed to download a distribution for the running SDK. '
        'Expected either one of %s to be found in the download folder.' %
        (expected_files))