#!/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())