# coding=utf-8
# Copyright 2015 Foursquare Labs Inc. All Rights Reserved.

from __future__ import absolute_import, division, print_function, unicode_literals

import os
import re
import shutil

from pants.backend.jvm.targets.scala_library import ScalaLibrary
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.build_graph.resources import Resources
from pants.option.custom_types import target_list_option, target_option
from pants.task.simple_codegen_task import SimpleCodegenTask
from pants.util.contextutil import temporary_dir
from pants.util.dirutil import safe_mkdir
from pants.util.memo import memoized_property

from fsqio.pants.spindle.targets.spindle_thrift_library import SpindleThriftLibrary
from fsqio.pants.spindle.tasks.spindle_task import SpindleTask


NAMESPACE_PARSER = re.compile(r'^\s*namespace\s+([^\s]+)\s+([^\s]+)\s*$')


class SpindleGen(SpindleTask, SimpleCodegenTask):
  """Generate codegen for spindle libraries."""
  # In order to ensure that any codegen generated by Spindle respects potential changes to the Spindle source code,
  # the SpindleGen task invokes a binary that is bootstrapped by an shelled pants run that is injected into the
  # task graph.
  # Practically, this means that if changes to the Spindle binary's source code are detected, a task is run
  # at the beginning of the execution that invokes pants in a subshell and compiles the binary. That SpindleBinary is
  # then used here during codegen, hopefully making the experience of working on Spindle a semi-reasonable experience.

  @classmethod
  def register_options(cls, register):
    super(SpindleGen, cls).register_options(register)
    register(
      '--runtime-dependency',
      advanced=True,
      fingerprint=True,
      type=target_list_option,
      help='A list of targets that all spindle codegen depends on at runtime.',
    )
    register(
      '--scala-ssp-template',
      fingerprint=True,
      advanced=True,
      type=target_option,
      help='Use this target as the scala templates for spindle codegen (required to be 1 target).',
    )
    register(
      '--write-annotations-json',
      fingerprint=True,
      advanced=True,
      type=bool,
      default=False,
      help='output *.annotations.json files for runtime class lists without reflection',
    )

  @classmethod
  def implementation_version(cls):
    return super(SpindleGen, cls).implementation_version() + [('SpindleGen', 2)]

  @classmethod
  def prepare(cls, options, round_manager):
    super(SpindleGen, cls).prepare(options, round_manager)
    round_manager.require('spindle_binary')

  @staticmethod
  def namespace_out(cache_dir):
    return os.path.join(cache_dir, 'scala_record')

  @staticmethod
  def scalate_workdir(cache_dir):
    return os.path.join(cache_dir, 'scalate')

  @memoized_property
  def scala_template(self):
    return self.get_ssp_templates(
      self.context.build_graph.get_target_from_spec(self.get_options().scala_ssp_template)
    )

  @memoized_property
  def _jvm_options(self):
    return self.get_options().jvm_options

  @memoized_property
  def _annotations(self):
    return self.get_options().write_annotations_json

  @memoized_property
  def spindle_binary(self):
    spindle_products = self.context.products.get('spindle_binary')
    products = spindle_products.get(self.spindle_target)
    if products:
      for directory, product in products.items():
        for filename in product:
          binary = os.path.join(directory, filename)
          if len(product) == 1 and os.path.isfile(binary):
            return binary
          # In this case we know there is just one product in the spindle_binary product.

    raise TaskError(
      'Spindle requires a single snapshot of Spindle runtime source code in order to bootstrap.\n'
      'Found: {}\n'.format(product)
    )

  @memoized_property
  def _resolved_runtime_deps(self):
    """Returns a twitter.common.collections.orderedset.OrderedSet."""
    # Cache resolution of runtime deps, since all spindle targets share them.
    return self.resolve_deps(self.get_options().runtime_dependency)

  def synthetic_target_extra_dependencies(self, target, target_workdir):
    """Bundle in spindle common and, optionally, json annotations."""
    if self._annotations:
      json_files = [path for path in os.listdir(target_workdir) if path.endswith('.json')]
      if json_files:
        return self._resolved_runtime_deps | \
          [self.make_json_resource(os.path.relpath(target_workdir, get_buildroot()), json_files)]
    return self._resolved_runtime_deps

  @staticmethod
  def synthetic_target_type(target):
    return ScalaLibrary

  @staticmethod
  def is_gentarget(target):
    return isinstance(target, SpindleThriftLibrary)

  def execute(self):
    # Spindle has two different outputs:
    #   1) an intermediate cache of compiled templates in the scalate_workdir
    #   2) its primary output, generated scala and java files.
    #
    # A regular scala_library target that depends on a Spindle target only needs the generated source files, (2).
    # This is tricky to map, because Spindle, like regular Thrift, supports includes. Operating over a target with
    # includes will generate source files for both the primary target as well as every 'included' target.
    # So invoking spindle on a per-target basis can result in the same source files being generated dozens of times.
    # That also includes regenerating the intermediate cache in the scalate_workdir, and is all extraordinarily slow.
    #
    # We get around this here by using a heuristic to map just the source files generated by the target. This is
    # much safer then the original algorithm, which esssentially ran spindle once, in a reused output dir that had no
    # concept of namespacing or caching and disbursed the files to downstream targets from there.
    # This temporary directory serves as a stable workspace and cache for the entire task. The output is
    # copied to the vt.results_dir on a per-target basis and treated like regular Pants artifact from there.

    with self.invalidated(
      self.codegen_targets(),
      invalidate_dependents=True,
      fingerprint_strategy=self.get_fingerprint_strategy(),
    ) as invalidation_check:
      with self.context.new_workunit(name='execute', labels=[WorkUnitLabel.MULTITOOL]):
        with temporary_dir() as workdir:
          for vt in invalidation_check.all_vts:
            if not vt.valid:
              if self._do_validate_sources_present(vt.target):
                self.execute_codegen(vt.target, vt.results_dir, workdir)
                self._handle_duplicate_sources(vt.target, vt.results_dir)
                # TODO(awinter): mv next few lines to self.custom_copy(target) and then use
                #   inherited SimpleCodegenTask.execute.
                ns_out = self.namespace_out(workdir)
                self.cache_generated_files(
                  self.calculate_generated_sources(vt.target, ns_out), ns_out, vt.results_dir,
                )
            vt.update()
            self._inject_synthetic_target(vt.target, vt.results_dir, vt.cache_key)

  def make_json_resource(self, dirname, sources):
    """Return synthetic Resources target that provides the json annotations for a generated ScalaLibrary"""
    if os.path.isabs(dirname) or any(os.path.isabs(s) for s in sources):
      raise TaskError("Abs paths here will cause cross-machine build invalidation")
    return self.context.add_new_target(
      address=Address(dirname, 'json'),
      target_type=Resources,
      sources=sorted(sources),
    )

  # Passing the intermediate 'workdir' here is a break with the upstream API. But the performance hit of regenerating
  # the scalalate workdir and every dependent spindle target was more then I could consider, especially when we are
  # able to get the correctness fixes we need from the 'calculate_generated_sources' hackery.
  def execute_codegen(self, target, target_workdir, workdir):
    tool_args = [
      '--template', self.scala_template,
      '--namespace_out', self.namespace_out(workdir),
      '--working_dir', self.scalate_workdir(workdir),
    ]
    bases = {tgt.target_base for tgt in target.closure() if self.is_gentarget(tgt)}
    tool_args.extend(['--thrift_include', ':'.join(bases)])
    if self._annotations:
      tool_args.extend(['--write_annotations_json', 'true'])
    args = tool_args + target.sources_relative_to_buildroot()

    self.context.log.debug('Executing: {} {}\n'.format(self.spindle_target.main, ' '.join(args)))
    result = self.runjava(
      classpath=[self.spindle_binary],
      jvm_options=self._jvm_options,
      main=self.spindle_target.main,
      args=args,
      workunit_name='spindle-codegen',
    )
    if result != 0:
      raise TaskError('Spindle codegen exited non-zero ({})'.format(result))

  def cache_generated_files(self, generated_files, src, dst):
    """Copy a list of paths between directory roots, creating subdirs as needed.

    Note: this isn't the same as artifact caching.
    """
    self.context.log.debug('Moving the following files to {}:\n {}'.format(dst, ' '.join(generated_files)))
    for gen_file in generated_files:
      safe_mkdir(os.path.join(dst, os.path.dirname(gen_file)))
      new_path = os.path.join(dst, gen_file)
      old_path = os.path.join(src, gen_file)
      shutil.copy2(old_path, new_path)

  def calculate_generated_sources(self, target, ns_out=None):
    generated_scala_sources = [
      '{0}.{1}'.format(source, 'scala')
      for source in self.sources_generated_by_target(target)
    ]

    # generate json.
    if ns_out is not None:
      gen_json = [
        '.'.join(source.split('/')) + '.json'
        for source in self.sources_generated_by_target(target)
      ]
      # note: json only created when annotations are nonempty, hence the need to check presence
      actual_json = filter(
        lambda f: ns_out and os.path.exists(os.path.join(ns_out, f)),
        gen_json
      )
    else:
      # note: this case is for SpindleStubsGen
      actual_json = []

    return generated_scala_sources + actual_json

  def sources_generated_by_target(self, target):
    return [
      relative_genned_source
      for thrift_source in target.sources_relative_to_buildroot()
      for relative_genned_source in self.calculate_genfiles(thrift_source)
    ]

  # Hacky way to figure out which files get generated from a particular thrift source.
  # TODO: This could be emitted by the codegen tool.
  # That would also allow us to easily support 1:many codegen.
  @staticmethod
  def calculate_genfiles(source, lang=None):
    gen_lang = lang if lang else 'java'
    abs_source = os.path.join(get_buildroot(), source)
    with open(abs_source, 'r') as thrift:
      lines = thrift.readlines()
    namespaces = {}
    for line in lines:
      match = NAMESPACE_PARSER.match(line)
      if match and match.group(1) == gen_lang:
        namespace = match.group(2)
        namespaces[gen_lang] = namespace
        # After we find gen_lang namespace we can stop reading the file.
        break

    def calculate_scala_record_genfiles(namespace, source):
      """Returns the generated file basenames, add the file extension to get the full path."""
      basepath = namespace.replace('.', '/')
      name = os.path.splitext(os.path.basename(source))[0]
      return [os.path.join(basepath, name)]

    namespace = namespaces.get(gen_lang)
    if not namespace:
      raise TaskError('No namespace provided in source: {}'.format(abs_source))
    return calculate_scala_record_genfiles(namespace, abs_source)
