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

import os
import re
import shutil
import subprocess
from typing import Dict, Optional

from twitter.common.collections import OrderedSet

from pants.backend.codegen.thrift.lib.thrift import Thrift
from pants.base.build_environment import get_buildroot
from pants.base.exceptions import TaskError
from pants.base.workunit import WorkUnitLabel
from pants.option.custom_types import target_option
from pants.task.simple_codegen_task import SimpleCodegenTask
from pants.util.memo import memoized_property


class ApacheThriftGenBase(SimpleCodegenTask):
    # The name of the thrift generator to use. Subclasses must set.
    # E.g., java, py (see `thrift -help` for all available generators).
    thrift_generator: Optional[str] = None

    # Subclasses may set their own default generator options.
    default_gen_options_map: Optional[Dict[str, Optional[str]]] = None

    @classmethod
    def register_options(cls, register):
        super().register_options(register)

        # NB: As of thrift 0.9.2 there is 1 warning that -strict promotes to an error - missing a
        # struct field id.  If an artifact was cached with strict off, we must re-gen with strict on
        # since this case may be present and need to generate a thrift compile error.
        register(
            "--strict",
            default=True,
            fingerprint=True,
            type=bool,
            help="Run thrift compiler with strict warnings.",
        )
        # The old --gen-options was string-typed, so we keep it that way for backwards compatibility,
        # and reluctantly use the clunky name --gen-options-map for the new, map-typed options.
        # TODO: Do a deprecation cycle to restore the old name.
        register(
            "--gen-options-map",
            type=dict,
            advanced=True,
            fingerprint=True,
            default=cls.default_gen_options_map,
            help="Use these options for the {} generator.".format(cls.thrift_generator),
        )
        register(
            "--deps",
            advanced=True,
            type=list,
            member_type=target_option,
            help="A list of address specs pointing to dependencies of thrift generated code.",
        )
        register(
            "--service-deps",
            advanced=True,
            type=list,
            member_type=target_option,
            help="A list of address specs pointing to dependencies of thrift generated service "
            "code.  If not supplied, then --deps will be used for service deps.",
        )

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

    def synthetic_target_extra_dependencies(self, target, target_workdir):
        for source in target.sources_relative_to_buildroot():
            if self._declares_service(os.path.join(get_buildroot(), source)):
                return self._service_deps
        return self._deps

    def execute_codegen(self, target, target_workdir):
        target_cmd = self._thrift_cmd[:]

        bases = OrderedSet(tgt.target_base for tgt in target.closure() if self.is_gentarget(tgt))
        for base in bases:
            target_cmd.extend(("-I", base))

        if hasattr(target, "compiler_args"):
            target_cmd.extend(list(target.compiler_args or []))

        target_cmd.extend(("-o", target_workdir))

        for source in target.sources_relative_to_buildroot():
            cmd = target_cmd[:]
            cmd.append(os.path.join(get_buildroot(), source))
            with self.context.new_workunit(
                name=source, labels=[WorkUnitLabel.TOOL], cmd=" ".join(cmd)
            ) as workunit:
                result = subprocess.call(
                    cmd, stdout=workunit.output("stdout"), stderr=workunit.output("stderr")
                )
                if result != 0:
                    raise TaskError(
                        "{} ... exited non-zero ({})".format(self._thrift_binary, result)
                    )

        # The thrift compiler generates sources to a gen-[lang] subdir of the `-o` argument.  We
        # relocate the generated sources to the root of the `target_workdir` so that our base class
        # maps them properly.
        gen_dir = os.path.join(target_workdir, "gen-{}".format(self.thrift_generator))
        for path in os.listdir(gen_dir):
            shutil.move(os.path.join(gen_dir, path), target_workdir)
        os.rmdir(gen_dir)

    @memoized_property
    def _thrift_binary(self):
        return Thrift.scoped_instance(self).select(context=self.context)

    @memoized_property
    def _deps(self):
        deps = self.get_options().deps
        return list(self.resolve_deps(deps))

    @memoized_property
    def _service_deps(self):
        service_deps = self.get_options().service_deps
        return list(self.resolve_deps(service_deps)) if service_deps else self._deps

    SERVICE_PARSER = re.compile(r"^\s*service\s+(?:[^\s{]+)")

    def _declares_service(self, source):
        with open(source, "r") as thrift:
            return any(line for line in thrift if self.SERVICE_PARSER.search(line))

    @memoized_property
    def _thrift_cmd(self):
        cmd = [self._thrift_binary]

        def opt_str(item):
            return item[0] if not item[1] else "{}={}".format(*item)

        gen_opts_map = self.get_options().gen_options_map or {}
        gen_opts = [opt_str(item) for item in gen_opts_map.items()]

        generator_spec = (
            "{}:{}".format(self.thrift_generator, ",".join(gen_opts))
            if gen_opts
            else self.thrift_generator
        )
        cmd.extend(("--gen", generator_spec))

        if self.get_options().strict:
            cmd.append("-strict")
        if self.get_options().level == "debug":
            cmd.append("-verbose")
        return cmd