# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Library functions for IR extraction."""

import os
import pathlib
import re
import shutil
import subprocess
import multiprocessing
import functools
import json
import logging

from typing import Dict, List, Optional

_UNSPECIFIED_OVERRIDE = ["<UNSPECIFIED>"]


# TODO(ml-compiler-opt): maybe we can also convert here the cmdline file,from a
# \0 - separated list of strings, to a \n one.
def should_include_module(cmdline: str, match_regexp: Optional[str]) -> bool:
    """Determine if the module should be included."""
    if match_regexp is None:
        return True
    lines = cmdline.split("\0")
    return any(len(re.findall(match_regexp, l)) for l in lines)


def get_thinlto_index(cmdline: str, basedir: str) -> Optional[str]:
    opts = cmdline.split("\0")
    for option in opts:
        if option.startswith("-fthinlto-index"):
            return os.path.join(basedir, option.split("=")[1])
    return None


class TrainingIRExtractor:
    """IR and command line extraction from an object file."""

    def __init__(self, obj_relative_path, output_base_dir, obj_base_dir=None):
        """Set up a TrainingIRExtractor.

        Args:
          obj_relative_path: relative path to the input object file. It will be also
            used to construct the absolute path of the output IR and cmd files, by
            appending it to output_base_dir.
          output_base_dir: the directory under which the output will be produced.
          obj_base_dir: the base directory for all the input object files.
        """
        self._obj_relative_path = obj_relative_path
        self._output_base_dir = output_base_dir
        self._obj_base_dir = obj_base_dir if obj_base_dir is not None else ""

    def obj_base_dir(self):
        return self._obj_base_dir

    def output_base_dir(self):
        return self._output_base_dir

    def relative_output_path(self):
        return self._obj_relative_path

    def input_obj(self):
        return os.path.join(self.obj_base_dir(), self._obj_relative_path)

    def lld_src_bc(self):
        # .3.import.bc is the suffix attached to post-merge-pre-opt ('postimport')
        # IR bitcode saved by lld. It is hardcoded into lld.
        return os.path.join(
            self._obj_base_dir, self._obj_relative_path + ".3.import.bc"
        )

    def lld_src_thinlto(self):
        return os.path.join(self._obj_base_dir, self._obj_relative_path + ".thinlto.bc")

    def dest_dir(self):
        return os.path.join(
            self.output_base_dir(), os.path.dirname(self._obj_relative_path)
        )

    def module_name(self):
        return os.path.basename(self._obj_relative_path)

    def cmd_file(self):
        return os.path.join(self.dest_dir(), self.module_name() + ".cmd")

    def bc_file(self):
        return os.path.join(self.dest_dir(), self.module_name() + ".bc")

    def thinlto_index_file(self):
        return os.path.join(self.dest_dir(), self.module_name() + ".thinlto.bc")

    def _get_extraction_cmd_command(
        self, llvm_objcopy_path: str, cmd_section_name: str
    ):
        """Get llvm-objcopy and process args to a produce a command string that,
        when invoked, will extract the cmd section info ths self.cmd_file() file.
        """
        return [
            llvm_objcopy_path,
            "--dump-section=" + cmd_section_name + "=" + self.cmd_file(),
            self.input_obj(),
            "/dev/null",
        ]

    def _get_extraction_bc_command(
        self, llvm_objcopy_path: str, bitcode_section_name: str
    ):
        """Gets llvm-objcopy and process args to produce a command string that,
        when invoked, will extract the bitcode section into the self.bc_file()
        file.
        """
        return [
            llvm_objcopy_path,
            "--dump-section=" + bitcode_section_name + "=" + self.bc_file(),
            self.input_obj(),
            "/dev/null",
        ]

    def _extract_clang_artifacts(
        self,
        llvm_objcopy_path: str,
        cmd_filter: str,
        is_thinlto: bool,
        cmd_section_name: str,
        bitcode_section_name: str,
    ) -> Optional[str]:
        """Run llvm-objcopy to extract the .bc and command line."""
        if not os.path.exists(self.input_obj()):
            logging.info("%s does not exist.", self.input_obj())
            return None
        os.makedirs(self.dest_dir(), exist_ok=True)
        try:
            subprocess.check_output(
                self._get_extraction_cmd_command(llvm_objcopy_path, cmd_section_name),
                stderr=subprocess.STDOUT,
                encoding="utf-8",
            )
            if cmd_filter is not None or is_thinlto:
                with open(self.cmd_file(), encoding="utf-8") as f:
                    lines = f.readlines()
                assert len(lines) == 1
                cmdline = lines[0]
                if not should_include_module(cmdline, cmd_filter):
                    logging.info(
                        "Excluding module %s because it does not match the filter",
                        self.input_obj(),
                    )
                    os.remove(self.cmd_file())
                    return None
                if is_thinlto:
                    index_file = get_thinlto_index(cmdline, self.obj_base_dir())
                    shutil.copy(index_file, self.thinlto_index_file())

            subprocess.check_output(
                self._get_extraction_bc_command(
                    llvm_objcopy_path, bitcode_section_name
                ),
                stderr=subprocess.STDOUT,
                encoding="utf-8",
            )
        except subprocess.CalledProcessError as e:
            # This may happen if  .o file was build from asm (.S source).
            logging.warning("%s was not processed: %s", self.input_obj(), e)
            logging.info(e.output)
            return None
        assert (
            os.path.exists(self.cmd_file())
            and os.path.exists(self.bc_file())
            and (not is_thinlto or os.path.exists(self.thinlto_index_file()))
        )
        return self.relative_output_path()

    def _extract_lld_artifacts(self) -> Optional[str]:
        """Extract the .bc file with ThinLTO index from an lld ThinLTO invocation."""
        if not os.path.exists(self.lld_src_bc()):
            logging.info("%s does not exist.", self.lld_src_bc())
            return None
        if not os.path.exists(self.lld_src_thinlto()):
            logging.info("%s does not exist.", self.lld_src_thinlto())
            return None
        os.makedirs(self.dest_dir(), exist_ok=True)

        # Copy over the files
        shutil.copy(self.lld_src_bc(), self.bc_file())
        shutil.copy(self.lld_src_thinlto(), self.thinlto_index_file())

        assert os.path.exists(self.bc_file())
        assert os.path.exists(self.thinlto_index_file())
        return self._obj_relative_path

    def extract(
        self,
        llvm_objcopy_path: Optional[str] = None,
        cmd_filter: Optional[str] = None,
        thinlto_build: Optional[str] = None,
        cmd_section_name: Optional[str] = ".llvmcmd",
        bitcode_section_name: Optional[str] = ".llvmbc",
    ) -> Optional[str]:
        if thinlto_build == "local":
            return self._extract_lld_artifacts()
        return self._extract_clang_artifacts(
            llvm_objcopy_path=llvm_objcopy_path,
            cmd_filter=cmd_filter,
            is_thinlto=thinlto_build == "distributed",
            cmd_section_name=cmd_section_name,
            bitcode_section_name=bitcode_section_name,
        )


def convert_compile_command_to_objectfile(
    command: Dict[str, str], output_dir: str
) -> Optional[TrainingIRExtractor]:
    obj_base_dir = command["directory"]
    if "arguments" in command:
        cmd_parts = command["arguments"]
    elif "command" in command:
        cmd_parts = command["command"].split()
    else:
        logging.info("compile_commands element has no command and arguments")
        return None

    try:
        obj_index = cmd_parts.index("-o") + 1
    except ValueError:
        # This could happen if there are non-clang commands in compile_commands.json
        logging.info("Command has no -o option: %s", " ".join(cmd_parts))
        return None
    obj_rel_path = cmd_parts[obj_index]
    # TODO(mtrofin): is the obj_base_dir correct for thinlto index bc files?
    return TrainingIRExtractor(
        obj_relative_path=obj_rel_path,
        output_base_dir=output_dir,
        obj_base_dir=obj_base_dir,
    )


def load_from_compile_commands(
    json_array: List[Dict[str, str]], output_dir: str
) -> List[TrainingIRExtractor]:
    objs = [
        convert_compile_command_to_objectfile(cmd, output_dir) for cmd in json_array
    ]
    # Filter out None, in case there were non-clang commands in the .json
    return [obj for obj in objs if obj is not None]


def load_from_lld_params(
    params_array: List[str], obj_base_dir: str, output_dir: str
) -> List[TrainingIRExtractor]:
    """Create an ObjectFile array based on lld's parameters."""
    # yank out -o and the output. After that, anything not starting with '-', and
    # ending in a '.o', is an object file.
    try:
        minus_o_idx = params_array.index("-o")
        del params_array[minus_o_idx : minus_o_idx + 2]
        just_obj_paths = [
            o for o in params_array if not o.startswith("-") and o.endswith(".o")
        ]
    except ValueError:
        logging.info("This params file does not have an explicit -o option.")
        just_obj_paths = params_array

    def make_obj(obj_file: str) -> TrainingIRExtractor:
        return TrainingIRExtractor(
            obj_relative_path=obj_file,
            output_base_dir=output_dir,
            obj_base_dir=obj_base_dir,
        )

    return [make_obj(obj_file) for obj_file in just_obj_paths]


def load_from_directory(
    obj_base_dir: str, output_dir: str
) -> List[TrainingIRExtractor]:
    """Create an object file array by globbing an entire drectory.

    Args:
      obj_base_dir: The base build directory that all object files will be
        written out as being relative to.
      output_dir: The output directory where extracted .bc and .cmd files should
        be placed.
    """
    paths = [str(p) for p in pathlib.Path(obj_base_dir).glob("**/*.o")]

    def make_spec(obj_file: str):
        return TrainingIRExtractor(
            obj_relative_path=os.path.relpath(obj_file, start=obj_base_dir),
            output_base_dir=output_dir,
            obj_base_dir=obj_base_dir,
        )

    return [make_spec(path) for path in paths]


def load_for_lld_thinlto(
    obj_base_dir: str, output_dir: str
) -> List[TrainingIRExtractor]:
    # .3.import.bc is the suffix attached to post-merge-pre-opt ('postimport')
    # IR bitcode saved by lld. It is hardcoded into lld. ThinLTO index files
    # are also emitted next to the postimport bitcode, with the suffix
    # .thinlto.bc instead
    paths = [str(p) for p in pathlib.Path(obj_base_dir).glob("**/*.3.import.bc")]

    def make_spec(obj_file: str):
        return TrainingIRExtractor(
            # Cut away .3.import.bc
            obj_relative_path=os.path.relpath(obj_file, start=obj_base_dir)[:-12],
            output_base_dir=output_dir,
            obj_base_dir=obj_base_dir,
        )

    return [make_spec(path) for path in paths]


def load_bazel_aquery(aquery_json, obj_base_dir: str, output_dir: str):
    """Creates an object file array by looking at the JSON output of bazel aquery.

    Args:
      aquery_json: The JSON-formatted output of the bazel aquery command for
        the target of interest. The bazel aquery JSON should be a JSON
        serialized version of the analysis.ActionGraphContainer proto.
        https://github.com/bazelbuild/bazel/blob/master/src/main/protobuf/analysis_v2.proto
      obj_base_dir: The base build directory that all object files will be
        written out as arelative to.
      output_dir: The output directory where extracted .bc and .cmd files should
        be placed.
    """
    linker_params = []

    for action_info in aquery_json["actions"]:
        if action_info["mnemonic"] != "CppLink":
            continue
        linker_params = action_info["arguments"]

    return load_from_lld_params(linker_params, obj_base_dir, output_dir)


def run_extraction(
    objs: List[TrainingIRExtractor],
    num_workers: int,
    llvm_objcopy_path: str,
    cmd_filter: str,
    thinlto_build: str,
    cmd_section_name: str,
    bitcode_section_name: str,
):
    """Extracts all specified object files into the corpus directory.

    Args:
      objs: A list of TrainingIRExtractor Objects that represent the object files
        to extract bitcode/commands from.
      num_workers: The number of parallel processes to spawn to run the
        extraction.
      llvm_objcopy_path: The path to the llvm-objcopy to use for dumping sections.
      cmd_filter: A regular expression that is used to select for compilations
        performed with specific flags. If you want to include all compilations,
        set this to None.
      thinlto_build: Whether or not this is a ThinLTO build, and if so, the type.
        Set this to None if the build was not done with ThinLTO.
      cmd_section_name: The name of the command line section created by the
        bitcode embedding.
      bitcode_section_name: The name of the bitcode section created by the
        bitcode embedding.
    """
    extract_artifacts = functools.partial(
        TrainingIRExtractor.extract,
        llvm_objcopy_path=llvm_objcopy_path,
        cmd_filter=cmd_filter,
        thinlto_build=thinlto_build,
        cmd_section_name=cmd_section_name,
        bitcode_section_name=bitcode_section_name,
    )

    with multiprocessing.Pool(num_workers) as pool:
        relative_output_paths = pool.map(extract_artifacts, objs)
        pool.close()
        pool.join()
    return relative_output_paths


def write_corpus_manifest(
    thinlto_build: str, relative_output_paths: List[str], output_dir: str
):
    """Writes a corpus_manifest.json containing all necessary information about
    the corpus.

    Args:
      thinlto_build: Whether or not the build was done with ThinLTO and if so,
        what kind of ThinLTO. Set this to none if the build was not performed with
        ThinLTO.
      relative_output_paths: The relative (to the corpus directory) output paths
        of all the bitcode files that should be placed in the corpus manifest
      output_dir: The corpus directory where the corpus manifest should be
        placed.
    """
    # This comes first rather than later so global_command_override is at the top
    # of the .json after being written
    if thinlto_build == "local":
        corpus_description = {"global_command_override": _UNSPECIFIED_OVERRIDE}
    else:
        corpus_description = {}

    corpus_description.update(
        {
            "has_thinlto": thinlto_build is not None,
            "modules": [path for path in relative_output_paths if path is not None],
        }
    )

    with open(
        os.path.join(output_dir, "corpus_description.json"), "w", encoding="utf-8"
    ) as f:
        json.dump(corpus_description, f, indent=2)