#! /usr/bin/env python3
# vim: tabstop=4 softtabstop=4

"""
This script helps to generate code coverage report.

It uses Clang Source-based Code Coverage, see
https://clang.llvm.org/docs/SourceBasedCodeCoverage.html

Example usage:
    python3 scripts/code_coverage/coverage.py \\
        crypto_unittests url_unittests \\
        -b out/coverage -o out/report \\
        -f url/ -f crypto/

The command above runs the crypto_unittests and url_unittests targets and
creates a coverage report for them. The coverage report is filtered to include
only files and sub-directories under url/ and crypto/ directories.

For more options, please refer to tools/code_coverage/coverage.py -h.

This script is based on Chromium code coverage scripts. For an overview of how
code coverage works in Chromium, please refer to
https://chromium.googlesource.com/chromium/src/+/master/docs/testing/code_coverage.md
"""

import argparse
import json
import logging
import os
import shutil
import subprocess

import coverage_utils

# Absolute path to the root of the checkout.
SRC_ROOT_PATH = None

# Build directory, the value is parsed from command line arguments.
BUILD_DIR = None

# Output directory for generated artifacts, the value is parsed from command
# line arguemnts.
OUTPUT_DIR = None

# Name of the file extension for profraw data files.
PROFRAW_FILE_EXTENSION = "profraw"

# Name of the final profdata file, and this file needs to be passed to
# "llvm-cov" command in order to call "llvm-cov show" to inspect the
# line-by-line coverage of specific files.
PROFDATA_FILE_NAME = "coverage.profdata"

# Name of the file with summary information generated by llvm-cov export.
SUMMARY_FILE_NAME = "summary.json"

LOGS_DIR_NAME = "logs"

# Retry failed merges.
MERGE_RETRIES = 3

# String to replace with actual llvm profile path.
LLVM_PROFILE_FILE_PATH_SUBSTITUTION = "<llvm_profile_file_path>"


def generate_per_file_line_by_line_coverage_in_html(binary_paths, profdata_file_path, filters, ignore_filename_regex):
    """Generates per file line-by-line coverage in html using `llvm-cov show`.

    For a file with absolute path /a/b/x.cc, a html report is generated as:
    OUTPUT_DIR/coverage/a/b/x.cc.html. An index html file is also generated as:
    OUTPUT_DIR/index.html.

    Args:
        binary_paths: A list of paths to the instrumented binaries.
        profdata_file_path: A path to the profdata file.
        filters: A list of directories and files to get coverage for.
    """
    # llvm-cov show [options] -instr-profile PROFILE BIN [-object BIN,...]
    # [[-object BIN]] [SOURCES]
    # NOTE: For object files, the first one is specified as a positional argument,
    # and the rest are specified as keyword argument.
    logging.debug("Generating per file line by line coverage reports using "
                  "`llvm-cov show` command.")
    subprocess_cmd = [
        "llvm-cov", "show", "-format=html",
        "-output-dir={}".format(OUTPUT_DIR),
        "-instr-profile={}".format(profdata_file_path), binary_paths[0],
    ]
    subprocess_cmd.extend(
        ["-object=" + binary_path for binary_path in binary_paths[1:]])
    subprocess_cmd.extend(["-Xdemangler", "c++filt", "-Xdemangler", "-n"])
    subprocess_cmd.extend(filters)
    if ignore_filename_regex:
        subprocess_cmd.append("-ignore-filename-regex={}".format(ignore_filename_regex))
    subprocess.check_call(subprocess_cmd)
    logging.debug("Finished running `llvm-cov show` command.")


def get_logs_directory_path():
    """Path to the logs directory."""
    return os.path.join(coverage_utils.get_coverage_report_root_dir_path(OUTPUT_DIR), LOGS_DIR_NAME)


def get_profdata_file_path():
    """Path to the resulting .profdata file."""
    return os.path.join(coverage_utils.get_coverage_report_root_dir_path(OUTPUT_DIR), PROFDATA_FILE_NAME)


def get_summary_file_path():
    """The JSON file that contains coverage summary written by llvm-cov export."""
    return os.path.join(coverage_utils.get_coverage_report_root_dir_path(OUTPUT_DIR), SUMMARY_FILE_NAME)


def create_coverage_profdata_for_targets(targets, commands):
    """Runs target to generate the coverage profile data.

    Args:
        targets: A list of targets.
        commands: A list of commands used to run the targets.

    Returns:
        A relative path to the generated profdata file.
    """
    target_profdata_file_paths = get_target_profdata_paths_by_executing_commands(targets, commands)
    coverage_profdata_file_path = merge_target_profdata_files(target_profdata_file_paths)

    for target_profdata_file_path in target_profdata_file_paths:
        os.remove(target_profdata_file_path)

    return coverage_profdata_file_path


def get_target_profdata_paths_by_executing_commands(targets, commands):
    """Runs commands and returns the relative paths to the profraw data files.

    Args:
        targets: A list of targets built with coverage instrumentation.
        commands: A list of commands used to run the targets.

    Returns:
        A list of relative paths to the generated profraw data files.
    """
    logging.debug("Executing the test commands.")

    # Remove existing profraw data files.
    report_root_dir = coverage_utils.get_coverage_report_root_dir_path(OUTPUT_DIR)
    for file_or_dir in os.listdir(report_root_dir):
        if file_or_dir.endswith(PROFRAW_FILE_EXTENSION):
            os.remove(os.path.join(report_root_dir, file_or_dir))

    # Ensure that logs directory exists.
    if not os.path.exists(get_logs_directory_path()):
        os.makedirs(get_logs_directory_path())

    profdata_file_paths = []

    # Run all test targets to generate profraw data files.
    for target, command in zip(targets, commands):
        output_file_name = target + "_output.log"
        output_file_path = os.path.join(get_logs_directory_path(), output_file_name)

        profdata_file_path = None
        for _ in range(MERGE_RETRIES):
            logging.info("Running command `{}`, the output is redirected to '{}'.".format(command, output_file_path))

            # profraw files are generated inside the OUTPUT_DIR.
            output = execute_command(target, command, output_file_path)

            profraw_file_paths = []
            for file_or_dir in os.listdir(report_root_dir):
                if file_or_dir.endswith(PROFRAW_FILE_EXTENSION):
                    profraw_file_paths.append(os.path.join(report_root_dir, file_or_dir))

            assert profraw_file_paths, (
                "Running target '{}' failed to generate any profraw data file, "
                "please make sure the binary exists, is properly instrumented and "
                "does not crash.".format(target))

            assert isinstance(profraw_file_paths, list), (
                "Variable 'profraw_file_paths' is expected to be of type 'list', "
                "but it is a {}.".format(type(profraw_file_paths)))

            try:
                profdata_file_path = merge_target_profraw_files(target, profraw_file_paths)
                break
            except Exception:
                logging.info("Retrying...")
            finally:
                # Remove profraw files now so that they are not used in next iteration.
                for profraw_file_path in profraw_file_paths:
                    os.remove(profraw_file_path)

        assert profdata_file_path, (
            "Failed to merge target '{}' profraw files after {} retries.".format(target, MERGE_RETRIES))
        profdata_file_paths.append(profdata_file_path)

    logging.debug("Finished executing the test commands.")

    return profdata_file_paths


def get_environment_vars(profraw_file_path):
    """Return environment vars for subprocess, given a profraw file path."""
    env = os.environ.copy()
    env.update({
        "LLVM_PROFILE_FILE": profraw_file_path,
    })
    return env


def execute_command(target, command, output_file_path):
    """Runs a single command and generates a profraw data file."""
    # Per Clang "Source-based Code Coverage" doc:
    #
    # "%p" expands out to the process ID. It's not used by this scripts due to:
    # 1) If a target program spawns too many processess, it may exhaust all disk
    #    space available. For example, unit_tests writes thousands of .profraw
    #    files each of size 1GB+.
    # 2) If a target binary uses shared libraries, coverage profile data for them
    #    will be missing, resulting in incomplete coverage reports.
    #
    # "%Nm" expands out to the instrumented binary's signature. When this pattern
    # is specified, the runtime creates a pool of N raw profiles which are used
    # for on-line profile merging. The runtime takes care of selecting a raw
    # profile from the pool, locking it, and updating it before the program exits.
    # N must be between 1 and 9. The merge pool specifier can only occur once per
    # filename pattern.
    #
    # "%1m" is used when tests run in single process, such as fuzz targets.
    #
    # For other cases, "%4m" is chosen as it creates some level of parallelism,
    # but it's not too big to consume too much computing resource or disk space.
    profile_pattern_string = "%1m"
    if "mpi" in command:
        profile_pattern_string = "%4m"
    expected_profraw_file_name = os.extsep.join(
        [target, profile_pattern_string, PROFRAW_FILE_EXTENSION])
    expected_profraw_file_path = os.path.join(
        coverage_utils.get_coverage_report_root_dir_path(OUTPUT_DIR),
        expected_profraw_file_name)
    command = command.replace(LLVM_PROFILE_FILE_PATH_SUBSTITUTION,
                              expected_profraw_file_path)

    try:
        # Some fuzz targets or tests may write into stderr, redirect it as well.
        with open(output_file_path, "wb") as output_file_handle:
          subprocess.check_call(
              command,
              shell=True,
              stdout=output_file_handle,
              stderr=subprocess.STDOUT,
              env=get_environment_vars(expected_profraw_file_path))
    except subprocess.CalledProcessError as e:
        logging.warning("Command `{}` exited with non-zero return code.".format(command))

    return open(output_file_path, "rb").read()


def merge_target_profdata_files(profdata_file_paths):
    """Returns a relative path to coverage profdata file by merging target
    profdata files.

    Args:
        profdata_file_paths: A list of relative paths to the profdata data files
                             that are to be merged.

    Returns:
        A relative path to the merged coverage profdata file.

    Raises:
        CalledProcessError: An error occurred merging profdata files.
    """
    logging.debug("Merging target profraw files to create target profdata file.")
    profdata_file_path = get_profdata_file_path()
    try:
        subprocess_cmd = [
            "llvm-profdata", "merge", "-o", profdata_file_path, "-sparse=true"
        ]
        subprocess_cmd.extend(profdata_file_paths)

        output = subprocess.check_output(subprocess_cmd)
        logging.debug("Merge output: {}".format(output))
    except subprocess.CalledProcessError:
        logging.error("Failed to merge target profdata files to create coverage profdata.")
        raise

    logging.debug("Finished merging target profdata files.")
    logging.info("Code coverage profile data is created as '{}'.".format(profdata_file_path))
    return profdata_file_path


def merge_target_profraw_files(target, profraw_file_paths):
    """Returns a relative path to target profdata file by merging target
    profraw files.

    Args:
        profraw_file_paths: A list of relative paths to the profdata data files
                            that are to be merged.

    Returns:
        A relative path to the merged coverage profdata file.

    Raises:
        CalledProcessError: An error occurred merging profdata files.
    """
    logging.debug("Merging target profraw files to create target profdata file.")
    profdata_file_path = os.path.join(OUTPUT_DIR, "{}.profdata".format(target))

    try:
        subprocess_cmd = [
            "llvm-profdata", "merge", "-o", profdata_file_path, "-sparse=true"
        ]
        subprocess_cmd.extend(profraw_file_paths)

        output = subprocess.check_output(subprocess_cmd)
        logging.debug("Merge output: {}".format(output))
    except subprocess.CalledProcessError as error:
        logging.error("Failed to merge target profraw files to create target profdata.")
        raise error

    logging.debug("Finished merging target profraw files.")
    logging.info("Target '{}' profile data is created as '{}'.".format(target, profdata_file_path))
    return profdata_file_path


def generate_per_file_coverage_summary(binary_paths, profdata_file_path, filters, ignore_filename_regex):
    """Generates per file coverage summary using "llvm-cov export" command."""
    # llvm-cov export [options] -instr-profile PROFILE BIN [-object BIN,...]
    # [[-object BIN]] [SOURCES].
    # NOTE: For object files, the first one is specified as a positional argument,
    # and the rest are specified as keyword argument.
    logging.debug("Generating per-file code coverage summary using `llvm-cov export -summary-only` command.")
    subprocess_cmd = [
        "llvm-cov", "export", "-summary-only",
        "-instr-profile=" + profdata_file_path, binary_paths[0]
    ]
    subprocess_cmd.extend(
        ["-object=" + binary_path for binary_path in binary_paths[1:]])
    subprocess_cmd.extend(filters)
    if ignore_filename_regex:
        subprocess_cmd.append("-ignore-filename-regex={}".format(ignore_filename_regex))

    export_output = subprocess.check_output(subprocess_cmd)

    # Write output on the disk to be used by code coverage bot.
    with open(get_summary_file_path(), "wb") as f:
        f.write(export_output)

    return export_output


def verify_paths_and_return_absolutes(paths):
    """Verifies that the paths specified in |paths| exist and returns absolute versions.

    Args:
        paths: A list of files or directories.
    """
    absolute_paths = []
    for path in paths:
      absolute_path = os.path.join(SRC_ROOT_PATH, path)
      assert os.path.exists(absolute_path), "Path '{}' doesn't exist.".format(path)

      absolute_paths.append(absolute_path)

    return absolute_paths


def parse_targets_arguments(targets_args, build_dir):
    """Return binary paths from target names."""
    targets = []
    commands = []
    binary_paths = []
    for target in targets_args:
        if "::" in target:
           target, command = target.split("::")
           target = target.strip()
        else:
           command = target

        binary_path = os.path.join(build_dir, "bin", target)
        if not os.path.exists(binary_path):
            logging.warning("Target binary '{}' not found in build directory, skipping.".format(os.path.basename(binary_path)))
            continue

        targets.append(target)
        command = command.replace(target, binary_path)
        commands.append(command)
        binary_paths.append(binary_path)

    return targets, commands, binary_paths


def setup_output_dir():
    """Setup output directory."""
    if os.path.exists(OUTPUT_DIR):
        shutil.rmtree(OUTPUT_DIR)
    # Creates OUTPUT_DIR and its platform sub-directory.
    os.makedirs(coverage_utils.get_coverage_report_root_dir_path(OUTPUT_DIR))


def parse_command_arguments():
    """Adds and parses relevant arguments for tool comands.

    Returns:
      A dictionary representing the arguments.
    """
    arg_parser = argparse.ArgumentParser()
    arg_parser.usage = __doc__

    arg_parser.add_argument("-b", "--build-dir", type=str, required=True,
          help="The build directory, the path needs to be relative to the root of the checkout.")

    arg_parser.add_argument("-o", "--output-dir", type=str, required=True,
          help="Output directory for generated artifacts.")

    arg_parser.add_argument("-f", "--filters", action="append", required=False,
          help="Directories or files to get code coverage for, and all files under "
          "the directories are included recursively.")

    arg_parser.add_argument("-i", "--ignore-filename-regex", type=str,
          help="Skip source code files with file paths that match the given "
          "regular expression. For example, use -i=\'.*/out/.*|.*/third_party/.*\' "
          "to exclude files in third_party/ and out/ folders from the report.")

    arg_parser.add_argument("--no-file-view", action="store_true",
          help="Don't generate the file view in the coverage report. When there "
          "are large number of html files, the file view becomes heavy and may "
          "cause the browser to freeze, and this argument comes handy.")

    arg_parser.add_argument("-v", "--verbose", action="store_true",
          help="Prints additional output for diagnostics.")

    arg_parser.add_argument("-l", "--log_file", type=str,
          help="Redirects logs to a file.")

    arg_parser.add_argument("targets", nargs="+",
          help="The names of the test targets to run.")

    args = arg_parser.parse_args()
    return args


if __name__ == "__main__":
    """Execute tool commands."""

    # Change directory to source root to aid in relative paths calculations.
    SRC_ROOT_PATH = coverage_utils.get_full_path(
        os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
    os.chdir(SRC_ROOT_PATH)

    args = parse_command_arguments()
    coverage_utils.configure_logging(verbose=args.verbose, log_file=args.log_file)

    BUILD_DIR = coverage_utils.get_full_path(args.build_dir)
    OUTPUT_DIR = coverage_utils.get_full_path(args.output_dir)

    assert os.path.exists(BUILD_DIR), (
        "Build directory '{}' doesn't exist. Please rebuild TNL.".format(BUILD_DIR))

    absolute_filter_paths = []
    if args.filters:
        absolute_filter_paths = verify_paths_and_return_absolutes(args.filters)

    setup_output_dir()

    targets, commands, binary_paths = parse_targets_arguments(args.targets, args.build_dir)
    profdata_file_path = create_coverage_profdata_for_targets(targets, commands)

    logging.info("Generating code coverage report in html...")
    per_file_summary_data = generate_per_file_coverage_summary(
        binary_paths, profdata_file_path, absolute_filter_paths,
        args.ignore_filename_regex)
    generate_per_file_line_by_line_coverage_in_html(binary_paths, profdata_file_path,
                                             absolute_filter_paths,
                                             args.ignore_filename_regex)
    # Call prepare here.
    processor = coverage_utils.CoverageReportPostProcessor(
        OUTPUT_DIR,
        SRC_ROOT_PATH + "/src/TNL",
        per_file_summary_data,
        no_file_view=args.no_file_view)

    processor.prepare_html_report()
