#!/usr/bin/env python3
#
# Copyright 2013 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Instruments classes and jar files.

This script corresponds to the 'jacoco_instr' action in the Java build process.
Depending on whether jacoco_instrument is set, the 'jacoco_instr' action will
call the instrument command which accepts a jar and instruments it using
jacococli.jar.

"""

import argparse
import json
import os
import shutil
import sys
import zipfile

from util import build_utils
import action_helpers
import zip_helpers


# This should be same as recipe side token. See bit.ly/3STSPcE.
INSTRUMENT_ALL_JACOCO_OVERRIDE_TOKEN = 'INSTRUMENT_ALL_JACOCO'


def _AddArguments(parser):
    """Adds arguments related to instrumentation to parser.

    Args:
      parser: ArgumentParser object.
    """
    parser.add_argument(
        '--root-build-dir',
        required=True,
        help='Path to build directory rooted at checkout dir e.g. //out/Release',
    )
    parser.add_argument(
        '--input-path',
        required=True,
        help='Path to input file(s). Either the classes '
        'directory, or the path to a jar.',
    )
    parser.add_argument(
        '--output-path',
        required=True,
        help='Path to output final file(s) to. Either the '
        'final classes directory, or the directory in '
        'which to place the instrumented/copied jar.',
    )
    parser.add_argument(
        '--sources-json-file',
        required=True,
        help='File to create with the list of source directories '
        'and input path.',
    )
    parser.add_argument(
        '--target-sources-file',
        required=True,
        help='File containing newline-separated .java and .kt paths',
    )
    parser.add_argument(
        '--jacococli-jar', required=True, help='Path to jacococli.jar.'
    )
    parser.add_argument(
        '--files-to-instrument',
        help='Path to a file containing which source files are affected.',
    )


def _GetSourceDirsFromSourceFiles(source_files):
    """Returns list of directories for the files in |source_files|.

    Args:
      source_files: List of source files.

    Returns:
      List of source directories.
    """
    return list(
        set(os.path.dirname(source_file) for source_file in source_files)
    )


def _CreateSourcesJsonFile(
    source_dirs, input_path, sources_json_file, src_root, root_build_dir
):
    """Adds all normalized source directories and input path to
    |sources_json_file|.

    Args:
      source_dirs: List of source directories.
      input_path: The input path to non-instrumented class files.
      sources_json_file: File into which to write the list of source directories
        and input path.
      src_root: Root which sources added to the file should be relative to.
      root_build_dir: Build directory path rooted at checkout where
        sources_json_file is generated e.g. //out/Release

    Returns:
      An exit code.
    """
    src_root = os.path.abspath(src_root)
    relative_sources = []
    for s in source_dirs:
        abs_source = os.path.abspath(s)
        if abs_source[: len(src_root)] != src_root:
            print(
                'Error: found source directory not under repository root: %s %s'
                % (abs_source, src_root)
            )
            return 1
        rel_source = os.path.relpath(abs_source, src_root)

        relative_sources.append(rel_source)

    data = {}
    data['source_dirs'] = relative_sources
    data['input_path'] = []
    build_dir = os.path.join(src_root, root_build_dir[2:])
    data['output_dir'] = build_dir
    if input_path:
        data['input_path'].append(os.path.abspath(input_path))
    with open(sources_json_file, 'w', encoding='utf-8') as f:
        json.dump(data, f)
    return 0


def _GetAffectedClasses(jar_file, source_files):
    """Gets affected classes by affected source files to a jar.

    Args:
      jar_file: The jar file to get all members.
      source_files: The list of affected source files.

    Returns:
      A tuple of affected classes and unaffected members.
    """
    with zipfile.ZipFile(jar_file) as f:
        members = f.namelist()

    affected_classes = []
    unaffected_members = []

    for member in members:
        if not member.endswith('.class'):
            unaffected_members.append(member)
            continue

        is_affected = False
        index = member.find('$')
        if index == -1:
            index = member.find('.class')
        for source_file in source_files:
            if source_file.endswith(
                (member[:index] + '.java', member[:index] + '.kt')
            ):
                affected_classes.append(member)
                is_affected = True
                break
        if not is_affected:
            unaffected_members.append(member)

    return affected_classes, unaffected_members


def _InstrumentClassFiles(
    instrument_cmd,
    input_path,
    output_path,
    temp_dir,
    affected_source_files=None,
):
    """Instruments class files from input jar.

    Args:
      instrument_cmd: JaCoCo instrument command.
      input_path: The input path to non-instrumented jar.
      output_path: The output path to instrumented jar.
      temp_dir: The temporary directory.
      affected_source_files: The affected source file paths to input jar.
        Default is None, which means instrumenting all class files in jar.
    """
    affected_classes = None
    unaffected_members = None
    if affected_source_files:
        affected_classes, unaffected_members = _GetAffectedClasses(
            input_path, affected_source_files
        )

    # Extract affected class files.
    with zipfile.ZipFile(input_path) as f:
        f.extractall(temp_dir, affected_classes)

    instrumented_dir = os.path.join(temp_dir, 'instrumented')

    # Instrument extracted class files.
    instrument_cmd.extend([temp_dir, '--dest', instrumented_dir])
    build_utils.CheckOutput(instrument_cmd)

    if affected_source_files and unaffected_members:
        # Extract unaffected members to instrumented_dir.
        with zipfile.ZipFile(input_path) as f:
            f.extractall(instrumented_dir, unaffected_members)

    # Zip all files to output_path
    with action_helpers.atomic_output(output_path) as f:
        zip_helpers.zip_directory(f, instrumented_dir)


def _RunInstrumentCommand(parser):
    """Instruments class or Jar files using JaCoCo.

    Args:
      parser: ArgumentParser object.

    Returns:
      An exit code.
    """
    args = parser.parse_args()

    source_files = []
    if args.target_sources_file:
        source_files.extend(
            build_utils.ReadSourcesList(args.target_sources_file)
        )

    with build_utils.TempDir() as temp_dir:
        instrument_cmd = build_utils.JavaCmd() + [
            '-jar',
            args.jacococli_jar,
            'instrument',
        ]

        if not args.files_to_instrument:
            affected_source_files = None
        else:
            affected_files = build_utils.ReadSourcesList(
                args.files_to_instrument
            )
            # Check if coverage recipe decided to instrument everything by overriding
            # the try builder default setting(selective instrumentation). This can
            # happen in cases like a DEPS roll of jacoco library

            # Note: This token is preceded by ../../ because the paths to be
            # instrumented are expected to be relative to the build directory.
            # See _rebase_paths() at https://bit.ly/40oiixX
            token = '../../' + INSTRUMENT_ALL_JACOCO_OVERRIDE_TOKEN
            if token in affected_files:
                affected_source_files = None
            else:
                source_set = set(source_files)
                affected_source_files = [
                    f for f in affected_files if f in source_set
                ]

                # Copy input_path to output_path and return if no source file affected.
                if not affected_source_files:
                    shutil.copyfile(args.input_path, args.output_path)
                    # Create a dummy sources_json_file.
                    _CreateSourcesJsonFile(
                        [],
                        None,
                        args.sources_json_file,
                        build_utils.DIR_SOURCE_ROOT,
                        args.root_build_dir,
                    )
                    return 0
        _InstrumentClassFiles(
            instrument_cmd,
            args.input_path,
            args.output_path,
            temp_dir,
            affected_source_files,
        )

    source_dirs = _GetSourceDirsFromSourceFiles(source_files)
    # TODO(GYP): In GN, we are passed the list of sources, detecting source
    # directories, then walking them to re-establish the list of sources.
    # This can obviously be simplified!
    _CreateSourcesJsonFile(
        source_dirs,
        args.input_path,
        args.sources_json_file,
        build_utils.DIR_SOURCE_ROOT,
        args.root_build_dir,
    )

    return 0


def main():
    parser = argparse.ArgumentParser()
    _AddArguments(parser)
    _RunInstrumentCommand(parser)


if __name__ == '__main__':
    sys.exit(main())