# 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
import shutil

from pants.backend.jvm.targets.benchmark import Benchmark
from pants.backend.jvm.tasks.jvm_task import JvmTask
from pants.backend.jvm.tasks.jvm_tool_task_mixin import JvmToolTaskMixin
from pants.base.exceptions import TaskError
from pants.base.workunit import WorkUnitLabel
from pants.java.distribution.distribution import DistributionLocator
from pants.java.executor import SubprocessExecutor
from pants.java.jar.jar_dependency import JarDependency
from pants.java.util import execute_java


class BenchmarkRun(JvmToolTaskMixin, JvmTask):
  _CALIPER_MAIN = 'com.google.caliper.Runner'

  @classmethod
  def register_options(cls, register):
    super(BenchmarkRun, cls).register_options(register)
    register('--target', help='Name of the benchmark class. This is a mandatory argument.')
    register('--memory', type=bool, help='Enable memory profiling.')
    register('--debug', type=bool,
             help='Run the benchmark tool with in process debugging.')

    cls.register_jvm_tool(register,
                          'benchmark-tool',
                          classpath=[
                            # TODO (Eric Ayers) Caliper is old. Add jmh support?
                            # The caliper tool is shaded, and so shouldn't interfere with Guava 16.
                            JarDependency(org='com.google.caliper', name='caliper', rev='0.5-rc1'),
                          ],
                          classpath_spec='//:benchmark-caliper-0.5',
                          main=cls._CALIPER_MAIN)
    cls.register_jvm_tool(register,
                          'benchmark-agent',
                          classpath=[
                            JarDependency(org='com.google.code.java-allocation-instrumenter',
                                          name='java-allocation-instrumenter',
                                          rev='2.1',
                                          intransitive=True),
                          ],
                          classpath_spec='//:benchmark-java-allocation-instrumenter-2.1')

  @classmethod
  def subsystem_dependencies(cls):
    return super(BenchmarkRun, cls).subsystem_dependencies() + (DistributionLocator,)

  def __init__(self, *args, **kwargs):
    super(BenchmarkRun, self).__init__(*args, **kwargs)
    # TODO(Steve Gury):
    # Find all the target classes from the Benchmark target itself
    # https://jira.twitter.biz/browse/AWESOME-1938
    if not self.get_options().target:
      raise ValueError('Mandatory argument --target must be specified.')
    self.args.insert(0, self.get_options().target)
    if self.get_options().memory:
      self.args.append('--measureMemory')
    if self.get_options().debug:
      self.args.append('--debug')

  def _is_benchmark(self, target):
    return isinstance(target, Benchmark)

  def execute(self):
    targets = self.context.targets(predicate=self._is_benchmark)
    if not targets:
      raise TaskError('No jvm targets specified for benchmarking.')

    # For rewriting JDK classes to work, the JAR file has to be listed specifically in
    # the JAR manifest as something that goes in the bootclasspath.
    # The MANIFEST list a jar 'allocation.jar' this is why we have to rename it
    agent_tools_classpath = self.tool_classpath('benchmark-agent')
    agent_jar = agent_tools_classpath[0]
    allocation_jar = os.path.join(os.path.dirname(agent_jar), "allocation.jar")

    # TODO(Steve Gury): Find a solution to avoid copying the jar every run and being resilient
    # to version upgrade
    shutil.copyfile(agent_jar, allocation_jar)
    os.environ['ALLOCATION_JAR'] = str(allocation_jar)

    benchmark_tools_classpath = self.tool_classpath('benchmark-tool')

    # Collect a transitive classpath for the benchmark targets.
    classpath = self.classpath(targets, benchmark_tools_classpath)

    java_executor = SubprocessExecutor(DistributionLocator.cached())
    exit_code = execute_java(classpath=classpath,
                             main=self._CALIPER_MAIN,
                             jvm_options=self.jvm_options,
                             args=self.args,
                             workunit_factory=self.context.new_workunit,
                             workunit_name='caliper',
                             workunit_labels=[WorkUnitLabel.RUN],
                             executor=java_executor,
                             create_synthetic_jar=self.synthetic_classpath)
    if exit_code != 0:
      raise TaskError('java {} ... exited non-zero ({})'.format(self._CALIPER_MAIN, exit_code))
