# Copyright 2017 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

import os

from pants.backend.jvm.subsystems.shader import Shader
from pants.backend.jvm.tasks.nailgun_task import NailgunTask
from pants.base.exceptions import TaskError
from pants.base.revision import Revision
from pants.base.workunit import WorkUnitLabel

from pants.contrib.codeanalysis.tasks.indexable_java_targets import IndexableJavaTargets


class IndexJava(NailgunTask):
    cache_target_dirs = True

    _KYTHE_JAVA_INDEXER_MAIN = "com.google.devtools.kythe.analyzers.java.JavaIndexer"

    # The indexer unfortunately relies on some class-level global state (a JsonFormat.TypeRegistry,
    # see com.google.devtools.kythe.util.JsonUtil), and throws when it's initialized more than once.
    # An old fix for this (https://github.com/kythe/kythe/pull/117) no longer works. Until we can
    # sort that out, we force this task to run Java in a subprocess.
    execution_strategy = "subprocess"

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

    @classmethod
    def implementation_version(cls):
        # Bump this version to invalidate all past artifacts generated by this task.
        return super().implementation_version() + [
            ("IndexJava", 8),
        ]

    @classmethod
    def product_types(cls):
        return ["kythe_entries_files"]

    @classmethod
    def prepare(cls, options, round_manager):
        super().prepare(options, round_manager)
        round_manager.require_data("kzip_files")

    @classmethod
    def register_options(cls, register):
        super().register_options(register)
        cls.register_jvm_tool(register, "kythe-java-indexer", main=cls._KYTHE_JAVA_INDEXER_MAIN)
        cls.register_jvm_tool(
            register,
            "javac9",
            custom_rules=[Shader.exclude_package("com.sun", recursive=True),],
            main="com.sun.tools.javac.main.Main",
        )

    @staticmethod
    def _entries_file(vt):
        return os.path.join(vt.results_dir, "index.entries")

    def execute(self):
        indexable_targets = IndexableJavaTargets.global_instance().get(self.context)

        with self.invalidated(indexable_targets, invalidate_dependents=True) as invalidation_check:
            if invalidation_check.invalid_vts:
                indexer_cp = self.tool_classpath("kythe-java-indexer")
                jvm_options = []
                if self.dist.version < Revision.lenient("9"):
                    # When run on JDK8, Kythe requires javac9 on the bootclasspath.
                    javac9_cp = self.tool_classpath("javac9")
                    jvm_options.append("-Xbootclasspath/p:{}".format(":".join(javac9_cp)))
                jvm_options.extend(self.get_options().jvm_options)

                for vt in invalidation_check.invalid_vts:
                    self._index(vt, indexer_cp, jvm_options)

        for vt in invalidation_check.all_vts:
            entries = self._entries_file(vt)
            self.context.products.get_data("kythe_entries_files", dict)[vt.target] = entries

    def _index(self, vt, indexer_cp, jvm_options):
        self.context.log.info("Kythe indexing {}".format(vt.target.address.spec))
        kzip_file = self.context.products.get_data("kzip_files").get(vt.target)
        if not kzip_file:
            raise TaskError("No .kzip file found for {}".format(vt.target.address.spec))
        args = [kzip_file, "--emit_jvm", "semantic", "--out", self._entries_file(vt)]
        result = self.runjava(
            classpath=indexer_cp,
            main=self._KYTHE_JAVA_INDEXER_MAIN,
            jvm_options=jvm_options,
            args=args,
            workunit_name="kythe-index",
            workunit_labels=[WorkUnitLabel.COMPILER],
        )
        if result != 0:
            raise TaskError(
                "java {main} ... exited non-zero ({result})".format(
                    main=self._KYTHE_JAVA_INDEXER_MAIN, result=result
                )
            )