# ==================================================================================================
# Copyright 2014 Twitter, Inc.
# --------------------------------------------------------------------------------------------------
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this work except in compliance with the License.
# You may obtain a copy of the License in the LICENSE file, or 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.
# ==================================================================================================

from collections import defaultdict

import os

from twitter.common.contextutil import open_zip
from twitter.common.dirutil import safe_mkdir, safe_open

from pants.backend.core.targets.resources import Resources
from pants.backend.jvm.tasks.jvm_binary_task import JvmBinaryTask
from pants.base.address import SyntheticAddress
from pants.base.build_environment import get_buildroot
from pants.goal.products import MultipleRootedProducts


class Args(object):
  @staticmethod
  def maybe_has_args(target):
    return target.is_jvm

  @classmethod
  def for_target(cls, target, transitive, class_genmap):
    classnames = set()

    def add_classnames(tgt):
      if cls.maybe_has_args(tgt):
        class_mapping = class_genmap.get(tgt)
        if class_mapping:
          for base, classes in class_mapping.rel_paths():
            for clazz in classes:
              classnames.add(clazz.replace('.class', '').replace('/', '.'))

    if transitive:
      target.walk(add_classnames)
    else:
      add_classnames(target)

    return cls(classnames)

  def __init__(self, classnames):
    self._classnames = classnames

  def matches(self, line):
    line = line.strip()
    if not line:
      return False
    components = line.split(' ')
    keyname = components[0]
    if keyname in ('positional', 'field'):
      # Line format: [key] class field
      return components[1] in self._classnames
    elif keyname == 'parser':
      # Line format: [key] parsed-class parser-class
      return components[2] in self._classnames
    elif keyname == 'verifier':
      # Line format: [key] verified-class verification-annotation-class verifier-class
      return components[2] in self._classnames and components[3] in self._classnames
    else:
      # Unknown line (comments, ws, unknown configuration types)
      return True


class ResourceMapper(JvmBinaryTask):
  """Maps resource files generated by com.twitter.common#args-apt into a binary jar."""

  RESOURCE_RELDIR = 'com/twitter/common/args/apt'
  RESOURCE_BASENAME = 'cmdline.arg.info.txt'

  @classmethod
  def product_type(cls):
    return ['resources_by_target']

  def prepare(self, round_manager):
    round_manager.require_data('classes_by_target')

  def execute(self):
    # TODO(John Sirois): DRY - replace `lambda: defaultdict(MultipleRootedProducts)` with a simple
    # type. We should not need to know here that a dict is needed to hold the Target keys for
    # example.
    resources_by_target = self.context.products.get_data(
        'resources_by_target', lambda: defaultdict(MultipleRootedProducts))

    for target in self.context.targets(predicate=Args.maybe_has_args):
      self._add_args_resources(target, resources_by_target, transitive=False)

    for binary in self.context.targets(predicate=self.is_binary):
      self._add_args_resources(binary, resources_by_target, transitive=True)

  def _add_args_resources(self, target, resources_by_target, transitive):
    classes_by_target = self.context.products.get_data('classes_by_target')

    resource_dirs = set()

    def collect_resource_dirs(tgt):
      mapping = classes_by_target.get(tgt)
      if mapping:
        resource_dirs.update(os.path.join(base, self.RESOURCE_RELDIR)
                             for base, _ in mapping.rel_paths())

    if transitive:
      target.walk(collect_resource_dirs)
    else:
      collect_resource_dirs(target)

    lines = set()
    for resource_dir in resource_dirs:
      if os.path.exists(resource_dir):
        for file_name in os.listdir(resource_dir):
          if file_name.startswith(self.RESOURCE_BASENAME):
            with open(os.path.join(resource_dir, file_name)) as resource:
              lines.update(resource.readlines())

    if lines:
      args = Args.for_target(target, transitive, classes_by_target)
      lines = set(filter(args.matches, lines))

    if transitive:
      # Add args from any of our transitive external deps that have them.
      resource_path = os.path.join(self.RESOURCE_RELDIR, self.RESOURCE_BASENAME)
      for base_dir, jar_path in self.list_external_jar_dependencies(target):
        with open_zip(os.path.join(base_dir, jar_path)) as jar:
          for zipinfo in jar.infolist():
            if zipinfo.filename.startswith(resource_path):
              lines.update(jar.open(zipinfo).readlines())

    self._addargs(lines, target, resources_by_target, transitive)

  def _addargs(self, lines, target, resources_by_target, transitive):
    def is_configuration_info(line):
      line = line.strip()
      return line and not line.startswith('#')

    if any(filter(is_configuration_info, lines)):
      resource_root = os.path.join(self.workdir, target.id)
      safe_mkdir(resource_root)

      resource_index = '0' if transitive else '1'
      resource_rel_path = os.path.join(self.RESOURCE_RELDIR,
                                       '{0}.{1}'.format(self.RESOURCE_BASENAME, resource_index))

      content = '# Created by pants goal args-apt\n'
      content += ''.join(sorted(lines))

      resource_path = os.path.join(resource_root, resource_rel_path)
      with safe_open(resource_path, 'w') as resource_fp:
        resource_fp.write(content)
        self.context.log.debug(
            'Added args-apt resource file {rel_path} for {target}:\n{content}'
            .format(rel_path=os.path.relpath(resource_path, get_buildroot()),
                    target=target,
                    content=content))

      # We'll get 2 cmdline resources for binary targets - a 0 transitive and 1 immediate.
      # The 0 will always trump at java runtime, so its harmless to attach both resources in
      # separate dependencies to keep the implementation here simple; we just need distinct target
      # names for each as a result.
      target_name = '{0}.transitive'.format(target.id) if transitive else target.id

      spec_path = os.path.relpath(resource_root, get_buildroot())
      address = SyntheticAddress(spec_path=spec_path, target_name=target_name)
      args_apt_resources = self.context.add_new_target(address,
                                                       Resources,
                                                       derived_from=target,
                                                       sources=[resource_rel_path])
      target.inject_dependency(address)
      resources_by_target[args_apt_resources].add_rel_paths(resource_root, [resource_rel_path])
