#!/usr/bin/env vpython
#
# Copyright 2017 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.

import argparse
import json
import logging
import os
import stat
import shutil
import sys
import tempfile
import time

from webkitpy.common.system.filesystem import FileSystem
from webkitpy.common.system.log_utils import configure_logging
from webkitpy.layout_tests import merge_results


# The output JSON has the following arguments overwritten with a value from
# build properties. This occurs when '--build-properties' argument is provided
# and is mainly used when merging on build bots to provide better information
# about the build to the test results server.
# Format is a list of ('result json key', 'build property key').
RESULTS_JSON_VALUE_OVERRIDE_WITH_BUILD_PROPERTY = [
    ("build_number", "buildnumber"),
    ("builder_name", "buildername"),
    ("chromium_revision", "got_revision_cp"),
]


# ------------------------------------------------------------------------
def ensure_empty_dir(fs, directory, allow_existing, remove_existing):
    """Ensure an empty directory exists.

    Args:
        allow_existing (bool): Allow the empty directory to already exist.
        remove_existing (bool): Remove the contents if the directory
            already exists.
    """
    if not fs.exists(directory):
        fs.maybe_make_directory(directory)
        return

    logging.warning('Output directory exists %r', directory)
    if not allow_existing:
        raise IOError(
            ('Output directory %s exists!\n'
             'Use --allow-existing-output-directory to continue') % directory)

    if remove_existing and not fs.remove_contents(directory):
        raise IOError(
            ('Unable to remove output directory %s contents!\n'
             'See log output for errors.') % directory)


def main(argv):

    parser = argparse.ArgumentParser()
    parser.description = """\
Merges sharded layout test results into a single output directory.
"""
    parser.epilog = """\

If a post merge script is given, it will be run on the resulting merged output
directory. The script will be given the arguments plus
'--results_dir <output_directory>'.
"""

    parser.add_argument(
        '-v', '--verbose', action='store_true',
        help='Output information about merging progress.')

    parser.add_argument(
        '--results-json-override-value',
        nargs=2, metavar=('KEY', 'VALUE'), default=[],
        action='append',
        help='Override the value of a value in the result style JSON file '
             '(--result-jsons-override-value layout_test_dirs /tmp/output).')
    parser.add_argument(
        '--results-json-allow-unknown-if-matching',
        action='store_true', default=False,
        help='Allow unknown values in the result.json file as long as the '
             'value match on all shards.')

    parser.add_argument(
        '--output-directory',
        help='Directory to create the merged results in.')
    parser.add_argument(
        '--allow-existing-output-directory',
        action='store_true', default=False,
        help='Allow merging results into a directory which already exists.')
    parser.add_argument(
        '--remove-existing-output-directory',
        action='store_true', default=False,
        help='Remove merging results into a directory which already exists.')
    parser.add_argument(
        '--input-directories', nargs='+',
        help='Directories to merge the results from.')

    # Swarming Isolated Merge Script API
    # script.py \
    #     --build-properties /s/build.json \
    #     --output-json /tmp/output.json \
    #     --task-output-dir /path/to/task/output/dir \
    #     shard0/output.json \
    #     shard1/output.json
    parser.add_argument(
        '-o', '--output-json',
        help='(Swarming Isolated Merge Script API) Output JSON file to create.')
    parser.add_argument(
        '--build-properties',
        help='(Swarming Isolated Merge Script API) Build property JSON file provided by recipes.')
    parser.add_argument(
        '--task-output-dir',
        help='(Swarming Isolated Merge Script API) Directory containing all swarming task results.')
    parser.add_argument(
        '--results-json-override-with-build-property',
        nargs=2, metavar=('RESULT_JSON_KEY', 'BUILD_PROPERTY_KEY'), default=[],
        action='append',
        help='Override the value of a value in the result style JSON file '
             '(--result-jsons-override-value layout_test_dirs /tmp/output).')
    parser.add_argument(
        '--summary-json',
        help='(Swarming Isolated Merge Script API) Summary of shard state running on swarming.'
             '(Output of the swarming.py collect --task-summary-json=XXX command.)')

    # Script to run after merging the directories together. Normally used with archive_layout_test_results.py
    # scripts/slave/chromium/archive_layout_test_results.py \
    #     --results-dir /b/rr/tmpIcChUS/w/layout-test-results \
    #     --build-dir /b/rr/tmpIcChUS/w/src/out \
    #     --build-number 3665 \
    #     --builder-name 'WebKit Linux - RandomOrder' \
    #     --gs-bucket gs://chromium-layout-test-archives \
    #     --staging-dir /b/c/chrome_staging \
    #     --slave-utils-gsutil-py-path /b/rr/tmpIcChUS/rw/scripts/slave/.recipe_deps/depot_tools/gsutil.py
    # in dir /b/rr/tmpIcChUS/w
    parser.add_argument(
        '--post-merge-script',
        nargs='*',
        help='Script to call after the results have been merged.')

    # The position arguments depend on if we are using the isolated merge
    # script API mode or not.
    parser.add_argument(
        'positional', nargs='*',
        help='output.json from shards.')

    args = parser.parse_args(argv)
    if args.verbose:
        logging_level = logging.DEBUG
    else:
        logging_level = logging.INFO
    configure_logging(logging_level=logging_level)

    # Map the isolate arguments back to our output / input arguments.
    if args.output_json:
        logging.info('Running with isolated arguments')
        assert args.positional

        # TODO(tansell): Once removed everywhere, these lines can be removed.
        # For now we just check nobody is supply arguments we didn't expect.
        if args.results_json_override_with_build_property:
            for result_key, build_prop_key in args.results_json_override_with_build_property:
                assert (result_key, build_prop_key) in RESULTS_JSON_VALUE_OVERRIDE_WITH_BUILD_PROPERTY, (
                    "%s not in %s" % (result_key, RESULTS_JSON_VALUE_OVERRIDE_WITH_BUILD_PROPERTY))

        if not args.output_directory:
            args.output_directory = os.getcwd()
            args.allow_existing_output_directory = True
            args.remove_existing_output_directory = True

        assert not args.input_directories
        args.input_directories = [os.path.dirname(f) for f in args.positional]
        args.positional = []

    # Allow skipping the --input-directories bit, for example,
    #   merge-layout-test-results -o outputdir shard0 shard1 shard2
    if args.positional and not args.input_directories:
        args.input_directories = args.positional

    if not args.output_directory:
        args.output_directory = tempfile.mkdtemp(suffix='webkit_layout_test_results.')

    assert args.output_directory
    assert args.input_directories

    results_json_value_overrides = {}
    if args.build_properties:
        build_properties = json.loads(args.build_properties)

        for result_key, build_prop_key in RESULTS_JSON_VALUE_OVERRIDE_WITH_BUILD_PROPERTY:
            if build_prop_key not in build_properties:
                logging.warn('Required build property key "%s" was not found!', build_prop_key)
                continue
            results_json_value_overrides[result_key] = build_properties[build_prop_key]
        logging.debug('results_json_value_overrides: %r', results_json_value_overrides)

    merger = merge_results.LayoutTestDirMerger(
        results_json_value_overrides=results_json_value_overrides,
        results_json_allow_unknown_if_matching=args.results_json_allow_unknown_if_matching)

    ensure_empty_dir(
        FileSystem(),
        args.output_directory,
        allow_existing=args.allow_existing_output_directory,
        remove_existing=args.remove_existing_output_directory)

    merger.merge(args.output_directory, args.input_directories)

    merged_output_json = os.path.join(args.output_directory, 'output.json')
    if os.path.exists(merged_output_json) and args.output_json:
        logging.debug(
            'Copying output.json from %s to %s', merged_output_json, args.output_json)
        shutil.copyfile(merged_output_json, args.output_json)

    if args.post_merge_script:
        logging.debug('Changing directory to %s', args.output_directory)
        os.chdir(args.output_directory)

        post_script = list(args.post_merge_script)
        post_script.append('--result-dir', args.output_directory)

        logging.info('Running post merge script %r', post_script)
        os.execlp(post_script)

main(sys.argv[1:])
