#!/usr/bin/env vpython3

# 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.

"""Aggregates Jacoco coverage files to produce output."""


import argparse
import fnmatch
import json
import os
import sys

import devil_chromium
from devil.utils import cmd_helper
from pylib.constants import host_paths

# Source paths should be passed to Jacoco in a way that the relative file paths
# reflect the class package name.
_PARTIAL_PACKAGE_NAMES = ['com/google', 'org/chromium']

# The sources_json_file is generated by jacoco_instr.py with source directories
# and input path to non-instrumented jars.
# e.g.
# 'source_dirs': [
#   "chrome/android/java/src/org/chromium/chrome/browser/toolbar/bottom",
#   "chrome/android/java/src/org/chromium/chrome/browser/ui/system",
# ...]
# 'input_path':
#   '$CHROMIUM_OUTPUT_DIR/\
#    obj/chrome/android/features/tab_ui/java__process_prebuilt-filtered.jar'

_SOURCES_JSON_FILES_SUFFIX = '__jacoco_sources.json'


def _CreateClassfileArgs(class_files, report_type, include_substr=None):
  """Returns a filtered list of files with classfile option.

  Args:
    class_files: A list of class files.
    report_type: A string indicating if device or host files are desired.
    include_substr: A substring that must be present to include the file.

  Returns:
    A list of files that don't use the suffix.
  """
  # These should match the jar class files generated in internal_rules.gni
  search_jar_suffix = '%s.filter.jar' % report_type
  result_class_files = []
  for f in class_files:
    include_file = False
    if f.endswith(search_jar_suffix):
      include_file = True

    # If include_substr is specified, remove files that don't have the
    # required substring.
    if include_file and include_substr and include_substr not in f:
      include_file = False
    if include_file:
      result_class_files += ['--classfiles', f]

  return result_class_files


def _GenerateReportOutputArgs(args, class_files, report_type):
  cmd = _CreateClassfileArgs(class_files, report_type,
                             args.include_substr_filter)
  if args.format == 'html':
    report_dir = os.path.join(args.output_dir, report_type)
    if not os.path.exists(report_dir):
      os.makedirs(report_dir)
    cmd += ['--html', report_dir]
  elif args.format == 'xml':
    cmd += ['--xml', args.output_file]
  elif args.format == 'csv':
    cmd += ['--csv', args.output_file]

  return cmd


def _GetFilesWithSuffix(root_dir, suffix):
  """Gets all files with a given suffix.

  Args:
    root_dir: Directory in which to search for files.
    suffix: Suffix to look for.

  Returns:
    A list of absolute paths to files that match.
  """
  files = []
  for root, _, filenames in os.walk(root_dir):
    basenames = fnmatch.filter(filenames, '*' + suffix)
    files.extend([os.path.join(root, basename) for basename in basenames])

  return files


def _GetExecFiles(root_dir, exclude_substr=None):
  """ Gets all .exec files

  Args:
    root_dir: Root directory in which to search for files.
    exclude_substr: Substring which should be absent in filename. If None, all
      files are selected.

  Returns:
    A list of absolute paths to .exec files

  """
  all_exec_files = _GetFilesWithSuffix(root_dir, ".exec")
  valid_exec_files = []
  for exec_file in all_exec_files:
    if not exclude_substr or exclude_substr not in exec_file:
      valid_exec_files.append(exec_file)
  return valid_exec_files


def _ParseArguments(parser):
  """Parses the command line arguments.

  Args:
    parser: ArgumentParser object.

  Returns:
    The parsed arguments.
  """
  parser.add_argument(
      '--format',
      required=True,
      choices=['html', 'xml', 'csv'],
      help='Output report format. Choose one from html, xml and csv.')
  parser.add_argument(
      '--device-or-host',
      choices=['device', 'host'],
      help='Selection on whether to use the device classpath files or the '
      'host classpath files. Host would typically be used for junit tests '
      ' and device for tests that run on the device. Only used for xml and csv'
      ' reports.')
  parser.add_argument('--include-substr-filter',
                      help='Substring that must be included in classjars.',
                      type=str,
                      default='')
  parser.add_argument('--output-dir', help='html report output directory.')
  parser.add_argument('--output-file',
                      help='xml file to write device coverage results.')
  parser.add_argument(
      '--coverage-dir',
      required=True,
      help='Root of the directory in which to search for '
      'coverage data (.exec) files.')
  parser.add_argument('--exec-filename-excludes',
                      required=False,
                      help='Excludes .exec files which contain a particular '
                      'substring in their name')
  parser.add_argument(
      '--sources-json-dir',
      help='Root of the directory in which to search for '
      '*__jacoco_sources.json files.')
  parser.add_argument(
      '--class-files',
      nargs='+',
      help='Location of Java non-instrumented class files. '
      'Use non-instrumented jars instead of instrumented jars. '
      'e.g. use chrome_java__process_prebuilt_(host/device)_filter.jar instead'
      'of chrome_java__process_prebuilt-instrumented.jar')
  parser.add_argument(
      '--sources',
      nargs='+',
      help='Location of the source files. '
      'Specified source folders must be the direct parent of the folders '
      'that define the Java packages.'
      'e.g. <src_dir>/chrome/android/java/src/')
  parser.add_argument(
      '--cleanup',
      action='store_true',
      help='If set, removes coverage files generated at '
      'runtime.')
  args = parser.parse_args()

  if args.format == 'html' and not args.output_dir:
    parser.error('--output-dir needed for report.')
  if args.format in ('csv', 'xml'):
    if not args.output_file:
      parser.error('--output-file needed for xml/csv reports.')
    if not args.device_or_host and args.sources_json_dir:
      parser.error('--device-or-host selection needed with --sources-json-dir')
  if not (args.sources_json_dir or args.class_files):
    parser.error('At least either --sources-json-dir or --class-files needed.')
  return args


def main():
  parser = argparse.ArgumentParser()
  args = _ParseArguments(parser)

  devil_chromium.Initialize()

  coverage_files = _GetExecFiles(args.coverage_dir, args.exec_filename_excludes)
  if not coverage_files:
    parser.error('No coverage file found under %s' % args.coverage_dir)
  print('Found coverage files: %s' % str(coverage_files))

  class_files = []
  source_dirs = []
  if args.sources_json_dir:
    sources_json_files = _GetFilesWithSuffix(args.sources_json_dir,
                                             _SOURCES_JSON_FILES_SUFFIX)
    for f in sources_json_files:
      with open(f, 'r') as json_file:
        data = json.load(json_file)
        class_files.extend(data['input_path'])
        source_dirs.extend(data['source_dirs'])

  # Fix source directories as direct parent of Java packages.
  fixed_source_dirs = set()
  for path in source_dirs:
    for partial in _PARTIAL_PACKAGE_NAMES:
      if partial in path:
        fixed_dir = os.path.join(host_paths.DIR_SOURCE_ROOT,
                                 path[:path.index(partial)])
        fixed_source_dirs.add(fixed_dir)
        break

  if args.class_files:
    class_files += args.class_files
  if args.sources:
    fixed_source_dirs.update(args.sources)

  cmd = [
      'java', '-jar',
      os.path.join(host_paths.DIR_SOURCE_ROOT, 'third_party', 'jacoco', 'cipd',
                   'lib', 'jacococli.jar'), 'report'
  ] + coverage_files

  for source in fixed_source_dirs:
    cmd += ['--sourcefiles', source]

  if args.format == 'html':
    # Both reports are generated for html as the cq bot generates an html
    # report and we wouldn't know which one a developer needed.
    device_cmd = cmd + _GenerateReportOutputArgs(args, class_files, 'device')
    host_cmd = cmd + _GenerateReportOutputArgs(args, class_files, 'host')

    device_exit_code = cmd_helper.RunCmd(device_cmd)
    host_exit_code = cmd_helper.RunCmd(host_cmd)
    exit_code = device_exit_code or host_exit_code
  else:
    cmd = cmd + _GenerateReportOutputArgs(args, class_files,
                                          args.device_or_host)
    exit_code = cmd_helper.RunCmd(cmd)

  if args.cleanup:
    for f in coverage_files:
      os.remove(f)

  # Command tends to exit with status 0 when it actually failed.
  if not exit_code:
    if args.format == 'html':
      if not os.path.isdir(args.output_dir) or not os.listdir(args.output_dir):
        print('No report generated at %s' % args.output_dir)
        exit_code = 1
    elif not os.path.isfile(args.output_file):
      print('No device coverage report generated at %s' % args.output_file)
      exit_code = 1

  return exit_code


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