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

from __future__ import (absolute_import, division, generators, nested_scopes, print_function,
                        unicode_literals, with_statement)

import os
from collections import OrderedDict
from hashlib import sha1

from twitter.common.collections import OrderedSet

from pants.backend.codegen.protobuf.java.java_protobuf_library import JavaProtobufLibrary
from pants.backend.codegen.protobuf.subsystems.protoc import Protoc
from pants.backend.jvm.targets.jar_library import JarLibrary
from pants.backend.jvm.targets.java_library import JavaLibrary
from pants.backend.jvm.tasks.jar_import_products import JarImportProducts
from pants.base.build_environment import get_buildroot
from pants.base.exceptions import TaskError
from pants.base.workunit import WorkUnitLabel
from pants.build_graph.address import Address
from pants.fs.archive import ZIP
from pants.task.simple_codegen_task import SimpleCodegenTask
from pants.util.process_handler import subprocess


class ProtobufGen(SimpleCodegenTask):

  @classmethod
  def subsystem_dependencies(cls):
    return super(ProtobufGen, cls).subsystem_dependencies() + (Protoc.scoped(cls),)

  @classmethod
  def register_options(cls, register):
    super(ProtobufGen, cls).register_options(register)

    # The protoc plugin names are used as proxies for the identity of the protoc
    # executable environment here.  Plugin authors must include a version in the name for
    # proper invalidation of protobuf products in the face of plugin modification that affects
    # plugin outputs.
    register('--protoc-plugins', advanced=True, fingerprint=True, type=list,
             help='Names of protobuf plugins to invoke.  Protoc will look for an executable '
                  'named protoc-gen-$NAME on PATH.')
    register('--extra_path', advanced=True, type=list,
             help='Prepend this path onto PATH in the environment before executing protoc. '
                  'Intended to help protoc find its plugins.',
             default=None)
    register('--javadeps', advanced=True, type=list,
             help='Dependencies to bootstrap this task for generating java code.  When changing '
                  'this parameter you may also need to update --version.',
             default=['3rdparty:protobuf-java'])
    register('--import-from-root', type=bool, advanced=True,
             help='If set, add the buildroot to the path protoc searches for imports. '
                  'This enables using import paths relative to the build root in .proto files, '
                  'as recommended by the protoc documentation.')

  # TODO https://github.com/pantsbuild/pants/issues/604 prep start
  @classmethod
  def prepare(cls, options, round_manager):
    super(ProtobufGen, cls).prepare(options, round_manager)
    round_manager.require_data(JarImportProducts)
    round_manager.require_data('deferred_sources')
  # TODO https://github.com/pantsbuild/pants/issues/604 prep finish

  def __init__(self, *args, **kwargs):
    """Generates Java files from .proto files using the Google protobuf compiler."""
    super(ProtobufGen, self).__init__(*args, **kwargs)
    self.plugins = self.get_options().protoc_plugins or []
    self._extra_paths = self.get_options().extra_path or []

  @property
  def protobuf_binary(self):
    return Protoc.scoped_instance(self).select(context=self.context)

  @property
  def javadeps(self):
    return self.resolve_deps(self.get_options().javadeps or [])

  def synthetic_target_type(self, target):
    return JavaLibrary

  def synthetic_target_extra_dependencies(self, target, target_workdir):
    deps = OrderedSet()
    if target.imported_jars:
      # We need to add in the proto imports jars.
      jars_address = Address(os.path.relpath(target_workdir, get_buildroot()),
                             target.id + '-rjars')
      jars_target = self.context.add_new_target(jars_address,
                                                JarLibrary,
                                                jars=target.imported_jars,
                                                derived_from=target)
      deps.update([jars_target])
    deps.update(self.javadeps)
    return deps

  def is_gentarget(self, target):
    return isinstance(target, JavaProtobufLibrary)

  def execute_codegen(self, target, target_workdir):
    sources_by_base = self._calculate_sources(target)
    sources = target.sources_relative_to_buildroot()

    bases = OrderedSet()
    # Note that the root import must come first, otherwise protoc can get confused
    # when trying to resolve imports from the root against the import's source root.
    if self.get_options().import_from_root:
      bases.add('.')
    bases.update(sources_by_base.keys())
    bases.update(self._proto_path_imports([target]))

    gen_flag = '--java_out'

    gen = '{0}={1}'.format(gen_flag, target_workdir)

    args = [self.protobuf_binary, gen]

    if self.plugins:
      for plugin in self.plugins:
        args.append("--{0}_out={1}".format(plugin, target_workdir))

    for base in bases:
      args.append('--proto_path={0}'.format(base))

    args.extend(sources)

    # Tack on extra path entries. These can be used to find protoc plugins.
    protoc_environ = os.environ.copy()
    if self._extra_paths:
      protoc_environ['PATH'] = os.pathsep.join(self._extra_paths
                                               + protoc_environ['PATH'].split(os.pathsep))

    # Note: The test_source_ordering integration test scrapes this output, so modify it with care.
    self.context.log.debug('Executing: {0}'.format('\\\n  '.join(args)))
    with self.context.new_workunit(name='protoc',
                                   labels=[WorkUnitLabel.TOOL],
                                   cmd=' '.join(args)) as workunit:
      result = subprocess.call(args,
                               env=protoc_environ,
                               stdout=workunit.output('stdout'),
                               stderr=workunit.output('stderr'))
      if result != 0:
        raise TaskError('{} ... exited non-zero ({})'.format(self.protobuf_binary, result))

  def _calculate_sources(self, target):
    gentargets = OrderedSet()

    def add_to_gentargets(tgt):
      if self.is_gentarget(tgt):
        gentargets.add(tgt)

    self.context.build_graph.walk_transitive_dependency_graph(
      [target.address],
      add_to_gentargets,
      postorder=True)
    sources_by_base = OrderedDict()
    for target in gentargets:
      base = target.target_base
      if base not in sources_by_base:
        sources_by_base[base] = OrderedSet()
      sources_by_base[base].update(target.sources_relative_to_buildroot())
    return sources_by_base

  def _jars_to_directories(self, target):
    """Extracts and maps jars to directories containing their contents.

    :returns: a set of filepaths to directories containing the contents of jar.
    """
    files = set()
    jar_import_products = self.context.products.get_data(JarImportProducts)
    imports = jar_import_products.imports(target)
    for coordinate, jar in imports:
      files.add(self._extract_jar(coordinate, jar))
    return files

  def _extract_jar(self, coordinate, jar_path):
    """Extracts the jar to a subfolder of workdir/extracted and returns the path to it."""
    with open(jar_path, 'rb') as f:
      outdir = os.path.join(self.workdir, 'extracted', sha1(f.read()).hexdigest())
    if not os.path.exists(outdir):
      ZIP.extract(jar_path, outdir)
      self.context.log.debug('Extracting jar {jar} at {jar_path}.'
                             .format(jar=coordinate, jar_path=jar_path))
    else:
      self.context.log.debug('Jar {jar} already extracted at {jar_path}.'
                             .format(jar=coordinate, jar_path=jar_path))
    return outdir

  def _proto_path_imports(self, proto_targets):
    for target in proto_targets:
      for path in self._jars_to_directories(target):
        yield os.path.relpath(path, get_buildroot())
