#
# 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.
#

"""Generates Python proto modules and grpc stubs for Beam protos."""
import contextlib
import glob
import inspect
import logging
import multiprocessing
import os
import platform
import re
import shutil
import subprocess
import sys
import time
import warnings
from importlib import import_module

import pkg_resources

BEAM_PROTO_PATHS = [
    os.path.join('..', '..', 'model', 'pipeline', 'src', 'main', 'proto'),
    os.path.join('..', '..', 'model', 'job-management', 'src', 'main', 'proto'),
    os.path.join('..', '..', 'model', 'fn-execution', 'src', 'main', 'proto'),
    os.path.join('..', '..', 'model', 'interactive', 'src', 'main', 'proto'),
]

PYTHON_OUTPUT_PATH = os.path.join('apache_beam', 'portability', 'api')

MODEL_RESOURCES = [
    os.path.normpath('../../model/fn-execution/src/main/resources'\
            + '/org/apache/beam/model/fnexecution/v1/standard_coders.yaml'),
]


def generate_urn_files(log, out_dir):
  """
  Create python files with statically defined URN constants.

  Creates a <proto>_pb2_urn.py file for each <proto>_pb2.py file that contains
  an enum type.

  This works by importing each api.<proto>_pb2 module created by `protoc`,
  inspecting the module's contents, and generating a new side-car urn module.
  This is executed at build time rather than dynamically on import to ensure
  that it is compatible with static type checkers like mypy.
  """
  import google.protobuf.message as message
  import google.protobuf.pyext._message as pyext_message

  class Context(object):
    INDENT = '  '
    CAP_SPLIT = re.compile('([A-Z][^A-Z]*|^[a-z]+)')

    def __init__(self, indent=0):
      self.lines = []
      self.imports = set()
      self.empty_types = set()
      self._indent = indent

    @contextlib.contextmanager
    def indent(self):
      self._indent += 1
      yield
      self._indent -= 1

    def prepend(self, s):
      if s:
        self.lines.insert(0, (self.INDENT * self._indent) + s + '\n')
      else:
        self.lines.insert(0, '\n')

    def line(self, s):
      if s:
        self.lines.append((self.INDENT * self._indent) + s + '\n')
      else:
        self.lines.append('\n')

    def import_type(self, typ):
      modname = typ.__module__
      if modname in ('__builtin__', 'builtin'):
        return typ.__name__
      else:
        self.imports.add(modname)
        return modname + '.' + typ.__name__

    @staticmethod
    def is_message_type(obj):
      return isinstance(obj, type) and \
             issubclass(obj, message.Message)

    @staticmethod
    def is_enum_type(obj):
      return type(obj).__name__ == 'EnumTypeWrapper'

    def python_repr(self, obj):
      if isinstance(obj, message.Message):
        return self.message_repr(obj)
      elif isinstance(obj, (list,
                            pyext_message.RepeatedCompositeContainer,  # pylint: disable=c-extension-no-member
                            pyext_message.RepeatedScalarContainer)):  # pylint: disable=c-extension-no-member
        return '[%s]' % ', '.join(self.python_repr(x) for x in obj)
      else:
        return repr(obj)

    def empty_type(self, typ):
      name = ('EMPTY_' +
              '_'.join(x.upper()
                       for x in self.CAP_SPLIT.findall(typ.__name__)))
      self.empty_types.add('%s = %s()' % (name, self.import_type(typ)))
      return name

    def message_repr(self, msg):
      parts = []
      for field, value in msg.ListFields():
        parts.append('%s=%s' % (field.name, self.python_repr(value)))
      if parts:
        return '%s(%s)' % (self.import_type(type(msg)), ', '.join(parts))
      else:
        return self.empty_type(type(msg))

    def write_enum(self, enum_name, enum, indent):
      ctx = Context(indent=indent)

      with ctx.indent():
        for v in enum.DESCRIPTOR.values:
          extensions = v.GetOptions().Extensions

          prop = (
              extensions[beam_runner_api_pb2.beam_urn],
              extensions[beam_runner_api_pb2.beam_constant],
              extensions[metrics_pb2.monitoring_info_spec],
              extensions[metrics_pb2.label_props],
          )
          reprs = [self.python_repr(x) for x in prop]
          if all(x == "''" or x.startswith('EMPTY_') for x in reprs):
            continue
          ctx.line('%s = PropertiesFromEnumValue(%s)' %
                   (v.name, ', '.join(self.python_repr(x) for x in prop)))

      if ctx.lines:
        ctx.prepend('class %s(object):' % enum_name)
        ctx.prepend('')
        ctx.line('')
      return ctx.lines

    def write_message(self, message_name, message, indent=0):
      ctx = Context(indent=indent)

      with ctx.indent():
        for obj_name, obj in inspect.getmembers(message):
          if self.is_message_type(obj):
            ctx.lines += self.write_message(obj_name, obj, ctx._indent)
          elif self.is_enum_type(obj):
            ctx.lines += self.write_enum(obj_name, obj, ctx._indent)

      if ctx.lines:
        ctx.prepend('class %s(object):' % message_name)
        ctx.prepend('')
      return ctx.lines

  pb2_files = [path for path in glob.glob(os.path.join(out_dir, '*_pb2.py'))]
  api_path = os.path.dirname(pb2_files[0])
  sys.path.insert(0, os.path.dirname(api_path))

  def _import(m):
    return import_module('api.%s' % m)

  try:
    beam_runner_api_pb2 = _import('beam_runner_api_pb2')
    metrics_pb2 = _import('metrics_pb2')

    for pb2_file in pb2_files:
      modname = os.path.splitext(pb2_file)[0]
      out_file = modname + '_urns.py'
      modname = os.path.basename(modname)
      mod = _import(modname)

      ctx = Context()
      for obj_name, obj in inspect.getmembers(mod):
        if ctx.is_message_type(obj):
          ctx.lines += ctx.write_message(obj_name, obj)

      if ctx.lines:
        for line in reversed(sorted(ctx.empty_types)):
          ctx.prepend(line)

        for modname in reversed(sorted(ctx.imports)):
          ctx.prepend('from . import %s' % modname)

        ctx.prepend('from ..utils import PropertiesFromEnumValue')

        log.info("Writing urn stubs: %s" % out_file)
        with open(out_file, 'w') as f:
          f.writelines(ctx.lines)

  finally:
    sys.path.pop(0)


def _find_protoc_gen_mypy():
  # NOTE: this shouldn't be necessary if the virtualenv's environment
  #  is passed to tasks below it, since protoc will search the PATH itself
  fname = 'protoc-gen-mypy'
  if platform.system() == 'Windows':
    fname += ".exe"

  pathstr = os.environ.get('PATH')
  search_paths = pathstr.split(os.pathsep) if pathstr else []
  # should typically be installed into the venv's bin dir
  search_paths.insert(0, os.path.dirname(sys.executable))
  for path in search_paths:
    fullpath = os.path.join(path, fname)
    if os.path.exists(fullpath):
      return fullpath
  raise RuntimeError("Could not find %s in %s" %
                     (fname, ', '.join(search_paths)))


def generate_proto_files(force=False, log=None):

  try:
    import grpc_tools  # pylint: disable=unused-import
  except ImportError:
    warnings.warn('Installing grpcio-tools is recommended for development.')

  if log is None:
    logging.basicConfig()
    log = logging.getLogger(__name__)
    log.setLevel(logging.INFO)

  py_sdk_root = os.path.dirname(os.path.abspath(__file__))
  common = os.path.join(py_sdk_root, '..', 'common')
  proto_dirs = [os.path.join(py_sdk_root, path) for path in BEAM_PROTO_PATHS]
  proto_files = sum(
      [glob.glob(os.path.join(d, '*.proto')) for d in proto_dirs], [])
  out_dir = os.path.join(py_sdk_root, PYTHON_OUTPUT_PATH)
  out_files = [path for path in glob.glob(os.path.join(out_dir, '*_pb2.py'))]

  if out_files and not proto_files and not force:
    # We have out_files but no protos; assume they're up to date.
    # This is actually the common case (e.g. installation from an sdist).
    log.info('No proto files; using existing generated files.')
    return

  elif not out_files and not proto_files:
    if not os.path.exists(common):
      raise RuntimeError(
          'Not in apache git tree; unable to find proto definitions.')
    else:
      raise RuntimeError(
          'No proto files found in %s.' % proto_dirs)

  if force:
    regenerate = 'forced'
  elif not out_files:
    regenerate = 'no output files'
  elif len(out_files) < len(proto_files):
    regenerate = 'not enough output files'
  elif (
      min(os.path.getmtime(path) for path in out_files)
      <= max(os.path.getmtime(path)
             for path in proto_files + [os.path.realpath(__file__)])):
    regenerate = 'output files are out-of-date'
  elif len(out_files) > len(proto_files):
    regenerate = 'output files without corresponding .proto files'
    # too many output files: probably due to switching between git branches.
    # remove them so they don't trigger constant regeneration.
    for out_file in out_files:
      os.remove(out_file)
  else:
    regenerate = None

  if regenerate:
    try:
      from grpc_tools import protoc
    except ImportError:
      if platform.system() == 'Windows':
        # For Windows, grpcio-tools has to be installed manually.
        raise RuntimeError(
            'Cannot generate protos for Windows since grpcio-tools package is '
            'not installed. Please install this package manually '
            'using \'pip install grpcio-tools\'.')

      # Use a subprocess to avoid messing with this process' path and imports.
      # Note that this requires a separate module from setup.py for Windows:
      # https://docs.python.org/2/library/multiprocessing.html#windows
      p = multiprocessing.Process(
          target=_install_grpcio_tools_and_generate_proto_files,
          kwargs={'force': force})
      p.start()
      p.join()
      if p.exitcode:
        raise ValueError("Proto generation failed (see log for details).")
    else:
      log.info('Regenerating Python proto definitions (%s).' % regenerate)
      builtin_protos = pkg_resources.resource_filename('grpc_tools', '_proto')

      protoc_gen_mypy = _find_protoc_gen_mypy()

      log.info('Found protoc_gen_mypy at %s' % protoc_gen_mypy)

      args = (
          [sys.executable] +  # expecting to be called from command line
          ['--proto_path=%s' % builtin_protos] +
          ['--proto_path=%s' % d for d in proto_dirs] +
          ['--python_out=%s' % out_dir] +
          ['--plugin=protoc-gen-mypy=%s' % protoc_gen_mypy] +
          ['--mypy_out=%s' % out_dir] +
          # TODO(robertwb): Remove the prefix once it's the default.
          ['--grpc_python_out=grpc_2_0:%s' % out_dir] +
          proto_files)
      ret_code = protoc.main(args)
      if ret_code:
        raise RuntimeError(
            'Protoc returned non-zero status (see logs for details): '
            '%s' % ret_code)

      # copy resource files
      for path in MODEL_RESOURCES:
        shutil.copy2(os.path.join(py_sdk_root, path), out_dir)

      ret_code = subprocess.call(
          ["futurize", "--both-stages", "--write", "--no-diff", out_dir])
      if ret_code:
        raise RuntimeError(
            'Error applying futurize to generated protobuf python files.')

      generate_urn_files(log, out_dir)

  else:
    log.info('Skipping proto regeneration: all files up to date')


# Though wheels are available for grpcio-tools, setup_requires uses
# easy_install which doesn't understand them.  This means that it is
# compiled from scratch (which is expensive as it compiles the full
# protoc compiler).  Instead, we attempt to install a wheel in a temporary
# directory and add it to the path as needed.
# See https://github.com/pypa/setuptools/issues/377
def _install_grpcio_tools_and_generate_proto_files(force=False):
  py_sdk_root = os.path.dirname(os.path.abspath(__file__))
  install_path = os.path.join(py_sdk_root, '.eggs', 'grpcio-wheels')
  build_path = install_path + '-build'
  if os.path.exists(build_path):
    shutil.rmtree(build_path)
  logging.warning('Installing grpcio-tools into %s', install_path)
  try:
    start = time.time()
    subprocess.check_call(
        [sys.executable, '-m', 'pip', 'install',
         '--target', install_path, '--build', build_path,
         '--upgrade',
         '-r', os.path.join(py_sdk_root, 'build-requirements.txt')])
    logging.warning(
        'Installing grpcio-tools took %0.2f seconds.', time.time() - start)
  finally:
    sys.stderr.flush()
    shutil.rmtree(build_path, ignore_errors=True)
  sys.path.append(install_path)
  try:
    generate_proto_files(force=force)
  finally:
    sys.stderr.flush()


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  generate_proto_files(force=True)