# ==================================================================================================
# Copyright 2013 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.
# ==================================================================================================

import os
import textwrap

from contextlib import closing
from itertools import chain
from xml.etree import ElementTree

from twitter.common.collections import OrderedDict
from twitter.common.contextutil import open_zip as open_jar
from twitter.common.dirutil import  safe_open

from twitter.pants.base.build_environment import get_buildroot
from twitter.pants.base.hash_utils import hash_file
from twitter.pants.base.workunit import WorkUnit
from twitter.pants.tasks import TaskError


# Well known metadata file required to register scalac plugins with nsc.
_PLUGIN_INFO_FILE = 'scalac-plugin.xml'

class ZincUtils(object):
  """Convenient wrapper around zinc invocations.

  Instances are immutable, and all methods are reentrant (assuming that the java_runner is).
  """
  _ZINC_MAIN = 'com.typesafe.zinc.Main'

  def __init__(self, context, nailgun_task, jvm_options, color, jvm_tool_bootstrapper):
    self.context = context
    self._nailgun_task = nailgun_task  # We run zinc on this task's behalf.
    self._jvm_options = jvm_options
    self._color = color
    self._jvm_tool_bootstrapper = jvm_tool_bootstrapper

    # The target scala version.
    self._compile_bootstrap_key = 'scalac'
    compile_bootstrap_tools = context.config.getlist('scala-compile', 'compile-bootstrap-tools',
                                                     default=[':scala-compile-2.9.3'])
    self._jvm_tool_bootstrapper.register_jvm_tool(self._compile_bootstrap_key, compile_bootstrap_tools)

    # The zinc version (and the scala version it needs, which may differ from the target version).
    self._zinc_bootstrap_key = 'zinc'
    zinc_bootstrap_tools = context.config.getlist('scala-compile', 'zinc-bootstrap-tools', default=[':zinc'])
    self._jvm_tool_bootstrapper.register_jvm_tool(self._zinc_bootstrap_key, zinc_bootstrap_tools)

    # Compiler plugins.
    plugins_bootstrap_tools = context.config.getlist('scala-compile', 'scalac-plugin-bootstrap-tools',
                                                     default=[])
    if plugins_bootstrap_tools:
      self._plugins_bootstrap_key = 'plugins'
      self._jvm_tool_bootstrapper.register_jvm_tool(self._plugins_bootstrap_key, plugins_bootstrap_tools)
    else:
      self._plugins_bootstrap_key = None

  @property
  def _zinc_classpath(self):
    return self._jvm_tool_bootstrapper.get_jvm_tool_classpath(self._zinc_bootstrap_key)

  @property
  def _compiler_classpath(self):
    return self._jvm_tool_bootstrapper.get_jvm_tool_classpath(self._compile_bootstrap_key)

  @property
  def _plugin_jars(self):
    if self._plugins_bootstrap_key:
      return self._jvm_tool_bootstrapper.get_jvm_tool_classpath(self._plugins_bootstrap_key)
    else:
      return []

  @property
  def _zinc_jar_args(self):
    zinc_jars = ZincUtils.identify_zinc_jars(self._zinc_classpath)
    # The zinc jar names are also the flag names.
    return (list(chain.from_iterable([['-%s' % name, jarpath]
                                     for (name, jarpath) in zinc_jars.items()])) +
            ['-scala-path', ':'.join(self._compiler_classpath)])

  def _plugin_args(self):
    # Allow multiple flags and also comma-separated values in a single flag.
    if self.context.options.plugins is not None:
      plugin_names = [p for val in self.context.options.plugins for p in val.split(',')]
    else:
      plugin_names = self.context.config.getlist('scala-compile', 'scalac-plugins', default=[])

    plugin_args = self.context.config.getdict('scala-compile', 'scalac-plugin-args', default={})
    active_plugins = self.find_plugins(plugin_names)

    ret = []
    for name, jar in active_plugins.items():
      ret.append('-S-Xplugin:%s' % jar)
      for arg in plugin_args.get(name, []):
        ret.append('-S-P:%s:%s' % (name, arg))
    return ret

  def plugin_jars(self):
    """The jars containing code for enabled plugins."""
    return self._plugin_jars

  def _run_zinc(self, args, workunit_name='zinc', workunit_labels=None):
    zinc_args = [
      '-log-level', self.context.options.log_level or 'info',
    ]
    if not self._color:
      zinc_args.append('-no-color')
    zinc_args.extend(self._zinc_jar_args)
    zinc_args.extend(args)
    return self._nailgun_task.runjava(classpath=self._zinc_classpath,
                                      main=ZincUtils._ZINC_MAIN,
                                      jvm_options=self._jvm_options,
                                      args=zinc_args,
                                      workunit_name=workunit_name,
                                      workunit_labels=workunit_labels)

  def compile(self, opts, classpath, sources, output_dir, analysis_file, upstream_analysis_files):
    args = list(opts)  # Make a copy

    args.extend(self._plugin_args())

    if upstream_analysis_files:
      args.extend(
        ['-analysis-map', ','.join(['%s:%s' % kv for kv in upstream_analysis_files.items()])])

    args.extend([
      '-analysis-cache', analysis_file,
      # We add compiler_classpath to ensure the scala-library jar is on the classpath.
      # TODO: This also adds the compiler jar to the classpath, which compiled code shouldn't
      # usually need. Be more selective?
      '-classpath', ':'.join(self._compiler_classpath + classpath),
      '-d', output_dir
    ])
    args.extend(sources)
    self.log_zinc_file(analysis_file)
    if self._run_zinc(args, workunit_labels=[WorkUnit.COMPILER]):
      raise TaskError('Zinc compile failed.')

  @staticmethod
  def write_plugin_info(resources_dir, target):
    root = os.path.join(resources_dir, target.id)
    plugin_info_file = os.path.join(root, _PLUGIN_INFO_FILE)
    with safe_open(plugin_info_file, 'w') as f:
      f.write(textwrap.dedent('''
        <plugin>
          <name>%s</name>
          <classname>%s</classname>
        </plugin>
      ''' % (target.plugin, target.classname)).strip())
    return root, plugin_info_file

  # These are the names of the various jars zinc needs. They are, conveniently and
  # non-coincidentally, the names of the flags used to pass the jar locations to zinc.
  zinc_jar_names = ['compiler-interface', 'sbt-interface' ]

  @staticmethod
  def identify_zinc_jars(zinc_classpath):
    """Find the named jars in the zinc classpath.

    TODO: Make these mappings explicit instead of deriving them by jar name heuristics.
    """
    ret = OrderedDict()
    ret.update(ZincUtils.identify_jars(ZincUtils.zinc_jar_names, zinc_classpath))
    return ret

  @staticmethod
  def identify_jars(names, jars):
    jars_by_name = {}
    jars_and_filenames = [(x, os.path.basename(x)) for x in jars]

    for name in names:
      jar_for_name = None
      for jar, filename in jars_and_filenames:
        if filename.startswith(name):
          jar_for_name = jar
          break
      if jar_for_name is None:
        raise TaskError('Couldn\'t find jar named %s' % name)
      else:
        jars_by_name[name] = jar_for_name
    return jars_by_name

  def find_plugins(self, plugin_names):
    """Returns a map from plugin name to plugin jar."""
    plugin_names = set(plugin_names)
    plugins = {}
    buildroot = get_buildroot()
    # plugin_jars is the universe of all possible plugins and their transitive deps.
    # Here we select the ones to actually use.
    for jar in self.plugin_jars():
      with open_jar(jar, 'r') as jarfile:
        try:
          with closing(jarfile.open(_PLUGIN_INFO_FILE, 'r')) as plugin_info_file:
            plugin_info = ElementTree.parse(plugin_info_file).getroot()
          if plugin_info.tag != 'plugin':
            raise TaskError(
              'File %s in %s is not a valid scalac plugin descriptor' % (_PLUGIN_INFO_FILE, jar))
          name = plugin_info.find('name').text
          if name in plugin_names:
            if name in plugins:
              raise TaskError('Plugin %s defined in %s and in %s' % (name, plugins[name], jar))
            # It's important to use relative paths, as the compiler flags get embedded in the zinc
            # analysis file, and we port those between systems via the artifact cache.
            plugins[name] = os.path.relpath(jar, buildroot)
        except KeyError:
          pass

    unresolved_plugins = plugin_names - set(plugins.keys())
    if unresolved_plugins:
      raise TaskError('Could not find requested plugins: %s' % list(unresolved_plugins))
    return plugins

  def log_zinc_file(self, analysis_file):
    self.context.log.debug('Calling zinc on: %s (%s)' % (analysis_file, hash_file(analysis_file).upper() if os.path.exists(analysis_file) else 'nonexistent'))
