# Copyright 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# This file was copy-pasted over from:
# //build/scripts/slave/recipe_modules/swarming/resources/results_merger.py

# This file is responsbile for merging JSON test results in both the simplified
# JSON format and the Chromium JSON test results format version 3.

from __future__ import print_function

import copy
import json
import sys

# These fields must appear in the test result output
REQUIRED = {
    'interrupted',
    'num_failures_by_type',
    'seconds_since_epoch',
    'tests',
}

# These fields are optional, but must have the same value on all shards
OPTIONAL_MATCHING = (
    'builder_name',
    'build_number',
    'chromium_revision',
    'has_pretty_patch',
    'has_wdiff',
    'path_delimiter',
    'pixel_tests_enabled',
    'random_order_seed',
)

OPTIONAL_IGNORED = ('layout_tests_dir',)

# These fields are optional and will be summed together
OPTIONAL_COUNTS = (
    'fixable',
    'num_flaky',
    'num_passes',
    'num_regressions',
    'skipped',
    'skips',
)


class MergeException(Exception):
    pass


def merge_test_results(shard_results_list, test_cross_device=False):
    """ Merge list of results.

  Args:
    shard_results_list: list of results to merge. All the results must have the
      same format. Supported format are simplified JSON format & Chromium JSON
      test results format version 3 (see
      https://www.chromium.org/developers/the-json-test-results-format)
    test_cross_device: If true, some tests are running in multiple shards. This
      requires some extra handling on merging the values under 'tests'.

  Returns:
    a dictionary that represent the merged results. Its format follow the same
    format of all results in |shard_results_list|.
  """
    shard_results_list = [x for x in shard_results_list if x]
    if not shard_results_list:
        return {}

    if 'seconds_since_epoch' in shard_results_list[0]:
        return _merge_json_test_result_format(shard_results_list, test_cross_device)
    else:
        return _merge_simplified_json_format(shard_results_list)


def _merge_simplified_json_format(shard_results_list):
    # This code is specialized to the "simplified" JSON format that used to be
    # the standard for recipes.

    # These are the only keys we pay attention to in the output JSON.
    merged_results = {
        'successes': [],
        'failures': [],
        'valid': True,
    }

    for result_json in shard_results_list:
        successes = result_json.get('successes', [])
        failures = result_json.get('failures', [])
        valid = result_json.get('valid', True)

        if (not isinstance(successes, list) or not isinstance(failures, list) or
                not isinstance(valid, bool)):
            raise MergeException('Unexpected value type in %s' % result_json)  # pragma: no cover

        merged_results['successes'].extend(successes)
        merged_results['failures'].extend(failures)
        merged_results['valid'] = merged_results['valid'] and valid
    return merged_results


def _merge_json_test_result_format(shard_results_list, test_cross_device=False):
    # This code is specialized to the Chromium JSON test results format version 3:
    # https://www.chromium.org/developers/the-json-test-results-format

    # These are required fields for the JSON test result format version 3.
    merged_results = {
        'tests': {},
        'interrupted': False,
        'version': 3,
        'seconds_since_epoch': float('inf'),
        'num_failures_by_type': {}
    }

    # To make sure that we don't mutate existing shard_results_list.
    shard_results_list = copy.deepcopy(shard_results_list)
    for result_json in shard_results_list:
        # TODO(tansell): check whether this deepcopy is actually necessary.
        result_json = copy.deepcopy(result_json)

        # Check the version first
        version = result_json.pop('version', -1)
        if version != 3:
            raise MergeException(  # pragma: no cover (covered by
                # results_merger_unittest).
                'Unsupported version %s. Only version 3 is supported' % version)

        # Check the results for each shard have the required keys
        missing = REQUIRED - set(result_json)
        if missing:
            raise MergeException(  # pragma: no cover (covered by
                # results_merger_unittest).
                'Invalid json test results (missing %s)' % missing)

        # Curry merge_values for this result_json.
        # pylint: disable=cell-var-from-loop
        merge = lambda key, merge_func: merge_value(result_json, merged_results, key, merge_func)

        if test_cross_device:
            # Results from the same test(story) may be found on different
            # shards(devices). We need to handle the merging on story level.
            merge('tests', merge_tries_v2)
        else:
            # Traverse the result_json's test trie & merged_results's test tries in
            # DFS order & add the n to merged['tests'].
            merge('tests', merge_tries)

        # If any were interrupted, we are interrupted.
        merge('interrupted', lambda x, y: x | y)

        # Use the earliest seconds_since_epoch value
        merge('seconds_since_epoch', min)

        # Sum the number of failure types
        merge('num_failures_by_type', sum_dicts)

        # Optional values must match
        for optional_key in OPTIONAL_MATCHING:
            if optional_key not in result_json:
                continue

            if optional_key not in merged_results:
                # Set this value to None, then blindly copy over it.
                merged_results[optional_key] = None
                merge(optional_key, lambda src, dst: src)
            else:
                merge(optional_key, ensure_match)

        # Optional values ignored
        for optional_key in OPTIONAL_IGNORED:
            if optional_key in result_json:
                merged_results[optional_key] = result_json.pop(
                    # pragma: no cover (covered by
                    # results_merger_unittest).
                    optional_key)

        # Sum optional value counts
        for count_key in OPTIONAL_COUNTS:
            if count_key in result_json:  # pragma: no cover
                # TODO(mcgreevy): add coverage.
                merged_results.setdefault(count_key, 0)
                merge(count_key, lambda a, b: a + b)

        if result_json:
            raise MergeException(  # pragma: no cover (covered by
                # results_merger_unittest).
                'Unmergable values %s' % list(result_json.keys()))

    return merged_results


def merge_tries(source, dest):
    """ Merges test tries.

  This is intended for use as a merge_func parameter to merge_value.

  Args:
      source: A result json test trie.
      dest: A json test trie merge destination.
  """
    # merge_tries merges source into dest by performing a lock-step depth-first
    # traversal of dest and source.
    # pending_nodes contains a list of all sub-tries which have been reached but
    # need further merging.
    # Each element consists of a trie prefix, and a sub-trie from each of dest
    # and source which is reached via that prefix.
    pending_nodes = [('', dest, source)]
    while pending_nodes:
        prefix, dest_node, curr_node = pending_nodes.pop()
        for k, v in curr_node.items():
            if k in dest_node:
                if not isinstance(v, dict):
                    raise MergeException('%s:%s: %r not mergable, curr_node: %r\ndest_node: %r' %
                                         (prefix, k, v, curr_node, dest_node))
                pending_nodes.append(("%s:%s" % (prefix, k), dest_node[k], v))
            else:
                dest_node[k] = v
    return dest


def merge_tries_v2(source, dest):
    """ Merges test tries, and adds support for merging results for the same story
  from different devices, which is not supported on v1.

  This is intended for use as a merge_func parameter to merge_value.

  Args:
      source: A result json test trie.
      dest: A json test trie merge destination.
  """
    # merge_tries merges source into dest by performing a lock-step depth-first
    # traversal of dest and source.
    # pending_nodes contains a list of all sub-tries which have been reached but
    # need further merging.
    # Each element consists of a trie prefix, and a sub-trie from each of dest
    # and source which is reached via that prefix.
    pending_nodes = [('', dest, source)]
    while pending_nodes:
        prefix, dest_node, curr_node = pending_nodes.pop()
        for k, v in curr_node.items():
            if k in dest_node:
                if not isinstance(v, dict):
                    raise MergeException('%s:%s: %r not mergable, curr_node: %r\ndest_node: %r' %
                                         (prefix, k, v, curr_node, dest_node))
                elif 'actual' in v and 'expected' in v:
                    # v is test result of a story name which is already in dest
                    _merging_cross_device_results(v, dest_node[k])
                else:
                    pending_nodes.append(("%s:%s" % (prefix, k), dest_node[k], v))
            else:
                dest_node[k] = v
    return dest


def _merging_cross_device_results(src, dest):
    # 1. Merge the 'actual' field and update the is_unexpected based on new values
    dest['actual'] += ' %s' % src['actual']
    if any(actual != dest['expected'] for actual in dest['actual'].split()):
        dest['is_unexpected'] = True
    # 2. append each item under the 'artifacts' and 'times'.
    if 'artifacts' in src:
        if 'artifacts' in dest:
            for artifact, artifact_list in src['artifacts'].items():
                if artifact in dest['artifacts']:
                    dest['artifacts'][artifact] += artifact_list
                else:
                    dest['artifacts'][artifact] = artifact_list
        else:
            dest['artifacts'] = src['artifacts']
    if 'times' in src:
        if 'times' in dest:
            dest['times'] += src['times']
        else:
            dest['time'] = src['time']
            dest['times'] = src['times']
    # 3. remove the 'shard' because now the results are from multiple shards.
    if 'shard' in dest:
        del dest['shard']


def ensure_match(source, dest):
    """ Returns source if it matches dest.

  This is intended for use as a merge_func parameter to merge_value.

  Raises:
      MergeException if source != dest
  """
    if source != dest:
        raise MergeException(  # pragma: no cover (covered by
            # results_merger_unittest).
            "Values don't match: %s, %s" % (source, dest))
    return source


def sum_dicts(source, dest):
    """ Adds values from source to corresponding values in dest.

  This is intended for use as a merge_func parameter to merge_value.
  """
    for k, v in source.items():
        dest.setdefault(k, 0)
        dest[k] += v

    return dest


def merge_value(source, dest, key, merge_func):
    """ Merges a value from source to dest.

  The value is deleted from source.

  Args:
    source: A dictionary from which to pull a value, identified by key.
    dest: The dictionary into to which the value is to be merged.
    key: The key which identifies the value to be merged.
    merge_func(src, dst): A function which merges its src into dst,
        and returns the result. May modify dst. May raise a MergeException.

  Raises:
    MergeException if the values can not be merged.
  """
    try:
        dest[key] = merge_func(source[key], dest[key])
    except MergeException as e:
        e.message = "MergeFailure for %s\n%s" % (key, e.message)
        e.args = tuple([e.message] + list(e.args[1:]))
        raise
    del source[key]


def main(files):
    if len(files) < 2:
        sys.stderr.write("Not enough JSON files to merge.\n")
        return 1
    sys.stderr.write('Starting with %s\n' % files[0])
    result = json.load(open(files[0]))
    for f in files[1:]:
        sys.stderr.write('Merging %s\n' % f)
        result = merge_test_results([result, json.load(open(f))])
    print(json.dumps(result))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
