#!/usr/bin/env python
# Copyright 2019 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Merge results from code-coverage/pgo swarming runs.

This script merges code-coverage/pgo profiles from multiple shards. It also
merges the test results of the shards.

It is functionally similar to merge_steps.py but it accepts the parameters
passed by swarming api.
"""

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

import merge_lib as profile_merger


def _MergeAPIArgumentParser(*args, **kwargs):
  """Parameters passed to this merge script, as per:
  https://chromium.googlesource.com/chromium/tools/build/+/main/scripts/slave/recipe_modules/swarming/resources/merge_api.py
  """
  parser = argparse.ArgumentParser(*args, **kwargs)
  parser.add_argument('--build-properties',
                      help=argparse.SUPPRESS,
                      default='{}')
  parser.add_argument('--summary-json', help=argparse.SUPPRESS)
  parser.add_argument('--task-output-dir', help=argparse.SUPPRESS)
  parser.add_argument('-o',
                      '--output-json',
                      required=True,
                      help=argparse.SUPPRESS)
  parser.add_argument('jsons_to_merge', nargs='*', help=argparse.SUPPRESS)

  # Custom arguments for this merge script.
  parser.add_argument('--additional-merge-script',
                      help='additional merge script to run')
  parser.add_argument(
      '--additional-merge-script-args',
      help='JSON serialized string of args for the additional merge script')
  parser.add_argument('--profdata-dir',
                      required=True,
                      help='where to store the merged data')
  parser.add_argument('--llvm-profdata',
                      required=True,
                      help='path to llvm-profdata executable')
  parser.add_argument('--test-target-name', help='test target name')
  parser.add_argument('--java-coverage-dir',
                      help='directory for Java coverage data')
  parser.add_argument('--jacococli-path', help='path to jacococli.jar.')
  parser.add_argument(
      '--merged-jacoco-filename',
      help='filename used to uniquely name the merged exec file.')
  parser.add_argument('--javascript-coverage-dir',
                      help='directory for JavaScript coverage data')
  parser.add_argument('--chromium-src-dir',
                      help='directory for chromium/src checkout')
  parser.add_argument('--build-dir',
                      help='directory for the build directory in chromium/src')
  parser.add_argument(
      '--per-cl-coverage',
      action='store_true',
      help='set to indicate that this is a per-CL coverage build')
  parser.add_argument('--sparse',
                      action='store_true',
                      dest='sparse',
                      help='run llvm-profdata with the sparse flag.')
  # (crbug.com/1091310) - IR PGO is incompatible with the initial conversion
  # of .profraw -> .profdata that's run to detect validation errors.
  # Introducing a bypass flag that'll merge all .profraw directly to .profdata
  parser.add_argument(
      '--skip-validation',
      action='store_true',
      help='skip validation for good raw profile data. this will pass all '
      'raw profiles found to llvm-profdata to be merged. only applicable '
      'when input extension is .profraw.')
  return parser


def main():
  desc = 'Merge profraw files in <--task-output-dir> into a single profdata.'
  parser = _MergeAPIArgumentParser(description=desc)
  params = parser.parse_args()

  if params.java_coverage_dir:
    if not params.jacococli_path:
      parser.error('--jacococli-path required when merging Java coverage')
    if not params.merged_jacoco_filename:
      parser.error(
          '--merged-jacoco-filename required when merging Java coverage')

    output_path = os.path.join(params.java_coverage_dir,
                               '%s.exec' % params.merged_jacoco_filename)
    logging.info('Merging JaCoCo .exec files to %s', output_path)
    profile_merger.merge_java_exec_files(params.task_output_dir, output_path,
                                         params.jacococli_path)

  failed = False

  if params.javascript_coverage_dir and params.chromium_src_dir \
      and params.build_dir:
    current_dir = os.path.dirname(__file__)
    merge_js_results_script = os.path.join(current_dir, 'merge_js_results.py')
    args = [
        sys.executable,
        merge_js_results_script,
        '--task-output-dir',
        params.task_output_dir,
        '--javascript-coverage-dir',
        params.javascript_coverage_dir,
        '--chromium-src-dir',
        params.chromium_src_dir,
        '--build-dir',
        params.build_dir,
    ]

    rc = subprocess.call(args)
    if rc != 0:
      failed = True
      logging.warning('%s exited with %s', merge_js_results_script, rc)

  # The orderfile bot outputs only one orderfile that needs to be copied over
  # to the profdata dir, no additional merging necessary.
  if 'orderfile' in params.profdata_dir:
    logging.info('Detected orderfile bot in %s', params.profdata_dir)
    orderfiles = []
    for dir_path, _sub_dirs, file_names in os.walk(params.task_output_dir):
      logging.info('Found orderfile dir: %s', dir_path)
      for fn in file_names:
        if 'unpatched' in fn:
          logging.info('Ignoring %s as it is unpatched.', fn)
        else:
          logging.info('Found orderfile: %s', fn)
          orderfiles.append(os.path.join(dir_path, fn))
    assert len(orderfiles) == 1, f'More than one orderfile found: {orderfiles}'
    source_path = orderfiles[0]
    # It is important to copy the file to the root profile dir with a
    # predictable name for the bot to find it reliably.
    dest_path = os.path.join(os.path.dirname(params.profdata_dir),
                             'orderfile.out')
    logging.info('Copying orderfile from %s to %s', source_path, dest_path)
    shutil.copyfile(source_path, dest_path)

  # Name the output profdata file name as {test_target}.profdata or
  # default.profdata.
  output_prodata_filename = (params.test_target_name or 'default') + '.profdata'

  # NOTE: The profile data merge script must make sure that the profraw files
  # are deleted from the task output directory after merging, otherwise, other
  # test results merge script such as layout tests will treat them as json test
  # results files and result in errors.
  invalid_profiles, counter_overflows = profile_merger.merge_profiles(
      params.task_output_dir,
      os.path.join(params.profdata_dir, output_prodata_filename),
      '.profraw',
      params.llvm_profdata,
      sparse=params.sparse,
      skip_validation=params.skip_validation)

  # At the moment counter overflows overlap with invalid profiles, but this is
  # not guaranteed to remain the case indefinitely. To avoid future conflicts
  # treat these separately.
  if counter_overflows:
    with open(os.path.join(params.profdata_dir, 'profiles_with_overflows.json'),
              'w') as f:
      json.dump(counter_overflows, f)

  if invalid_profiles:
    with open(os.path.join(params.profdata_dir, 'invalid_profiles.json'),
              'w') as f:
      json.dump(invalid_profiles, f)

  # If given, always run the additional merge script, even if we only have one
  # output json. Merge scripts sometimes upload artifacts to cloud storage, or
  # do other processing which can be needed even if there's only one output.
  if params.additional_merge_script:
    new_args = [
        '--build-properties',
        params.build_properties,
        '--summary-json',
        params.summary_json,
        '--task-output-dir',
        params.task_output_dir,
        '--output-json',
        params.output_json,
    ]

    if params.additional_merge_script_args:
      new_args += json.loads(params.additional_merge_script_args)

    new_args += params.jsons_to_merge

    args = [sys.executable, params.additional_merge_script] + new_args
    rc = subprocess.call(args)
    if rc != 0:
      failed = True
      logging.warning('Additional merge script %s exited with %s',
                      params.additional_merge_script, rc)
  elif len(params.jsons_to_merge) == 1:
    logging.info('Only one output needs to be merged; directly copying it.')
    with open(params.jsons_to_merge[0]) as f_read:
      with open(params.output_json, 'w') as f_write:
        f_write.write(f_read.read())
  else:
    logging.warning(
        'This script was told to merge test results, but no additional merge '
        'script was given.')

  # TODO(crbug.com/40868908): Return non-zero if invalid_profiles is not None
  return 1 if failed else 0


if __name__ == '__main__':
  logging.basicConfig(format='[%(asctime)s %(levelname)s] %(message)s',
                      level=logging.INFO)
  sys.exit(main())