#!/usr/bin/env vpython
# Copyright 2018 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 multiprocessing
import sys
import textwrap

from core import benchmark_utils
from core import bot_platforms
from core import retrieve_story_timing
from core import sharding_map_generator

_SCRIPT_USAGE = """
Generate sharding maps for Telemetry benchmarks.

Every performance benchmark should be run on a same machine as long as possible
to preserve high fidelity of data monitoring. Hence in order to shard the
Telemetry benchmarks on multiple machines, we generate a JSON map that
specifies how benchmarks should be distributed on machines. There is one
sharding JSON map for every builder in the perf & perf.fyi waterfalls which are
specified by PerfPlatform classes in //tools/perf/core/bot_platforms.py.

Generating these JSON maps depends on how many Telemetry benchmarks
actually exist at the time. Because of this, CLs to generate the JSON maps
should never be automatically reverted, since the reverted state of the JSON map
files may not match with the true state of world.

"""


def GetParser():
  parser = argparse.ArgumentParser(
      description=_SCRIPT_USAGE, formatter_class=argparse.RawTextHelpFormatter)
  subparsers = parser.add_subparsers()

  parser_update = subparsers.add_parser('update')
  parser_update.add_argument(
      '--use-old-timing-data', '-o', action='store_true',
      help=('Whether to reuse existing builder timing data (stored in '
            '//tools/perf/core/shard_maps/timing_data/) and skip the step of '
            'fetching the most recent timing data from test results server. '
            'This flag is default to False. One typically uses this option '
            'when they need to fix the timing data to debug sharding '
            'generation.'),
      default=False)
  parser_update.add_argument(
      '--builders', '-b', action='append',
      help=('The builder names to reshard.'), default=[],
      choices=bot_platforms.ALL_PLATFORM_NAMES)
  parser_update.add_argument(
      '--waterfall', '-w', choices=['perf', 'perf-fyi', 'all'], default=None,
      help=('The name of waterfall whose builders to be resharded. If not '
            'specified, use all perf builders by default'))
  parser.add_argument(
      '--debug', action='store_true',
      help=('Whether to include detailed debug info of the sharding map in the '
            'shard maps.'), default=False)

  parser_update.set_defaults(func=_UpdateShardsForBuilders)

  parser_create = subparsers.add_parser('create')
  parser_create.add_argument(
      '--benchmark', help='The benchmark that you want to create shard for',
      required=True)
  parser_create.add_argument(
      '--timing-data-source', '-t', choices=bot_platforms.ALL_PLATFORM_NAMES,
      help='The timing data that you want to use. If not set, it will assume '
           'all stories use the same amount of time to run')
  parser_create.add_argument(
      # pinpoint typically has 16 machines for each hardware types, so we set
      # the default to use half of them to avoid starving the pool.
      '--shards-num', type=int, default=8,
      help="The number of shards you'd like to use, default is %(default)s")
  parser_create.add_argument(
      '--output-path', default='new_shard_map.json',
      help='Output file path for the shard map, default is `%(default)s`')
  parser_create.set_defaults(func=_CreateShardMapForBenchmark)
  return parser


def _GenerateBenchmarksToShardsList(benchmarks):
  """Return |benchmarks_to_shard| from given list of |benchmarks|.

    benchmarks_to_shard is a list all benchmarks to be sharded. Its
    structure is as follows:
    [{
       "name": "benchmark_1",
       "stories": [ "storyA", "storyB",...],
       "repeat": <number of pageset_repeat>
      },
      {
       "name": "benchmark_2",
       "stories": [ "storyA", "storyB",...],
       "repeat": <number of pageset_repeat>
      },
       ...
    ]

    The "stories" field contains a list of ordered story names. Notes that
    this should match the actual order of how the benchmark stories are
    executed for the sharding algorithm to be effective.
  """
  benchmarks_to_shard = []
  for b in benchmarks:
    benchmarks_to_shard.append({
        'name': b.Name(),
        'repeat': b().options.get('pageset_repeat', 1),
        'stories': benchmark_utils.GetBenchmarkStoryNames(b())
    })
  return benchmarks_to_shard


def _LoadTimingData(args):
  builder_name, timing_file_path = args
  data = retrieve_story_timing.FetchAverageStortyTimingData(
      configurations=[builder_name], num_last_days=5)
  with open(timing_file_path, 'w') as output_file:
    json.dump(data, output_file, indent=4, separators=(',', ': '))
  print 'Finish retrieve story timing data for %s' % repr(builder_name)


def _GenerateShardMap(
    builder, num_of_shards, output_path, debug, benchmark):
  timing_data = []
  if builder:
    with open(builder.timing_file_path) as f:
      timing_data = json.load(f)
  benchmarks_to_shard = _GenerateBenchmarksToShardsList(
      [b for b in builder.benchmarks_to_run if not benchmark or (
          b.Name() == benchmark)])
  sharding_map = sharding_map_generator.generate_sharding_map(
      benchmarks_to_shard, timing_data, num_shards=num_of_shards,
      debug=debug)
  with open(output_path, 'w') as output_file:
    json.dump(sharding_map, output_file, indent=4, separators=(',', ': '))


def _PromptWarning():
  message = ('This will regenerate the sharding maps for all perf benchmarks. '
             'Note that this will shuffle all the benchmarks on the shards, '
             'which can cause false regressions. In general this operation '
             'should only be done when the shards are too unbalanced or when '
             'benchmarks are added/removed. '
             'In addition, this a tricky operation and should '
             'only be done by Telemetry or Chrome Client Infrastructure '
             'team members. Upon landing the CL to update the shards maps, '
             'please notify Chromium perf sheriffs in '
             'perf-sheriffs@chromium.org and put a warning about expected '
             'false regressions in your CL '
             'description')
  print textwrap.fill(message, 70), '\n'
  answer = raw_input("Enter 'y' to continue: ")
  if answer != 'y':
    print 'Abort updating shard maps for benchmarks on perf waterfall'
    sys.exit(0)


def _UpdateShardsForBuilders(args):
  if args.builders:
    builders = {b for b in bot_platforms.ALL_PLATFORMS if b.name in
                args.builders}
  elif args.waterfall == 'perf':
    builders = bot_platforms.ALL_PERF_PLATFORMS
    _PromptWarning()
  elif args.waterfall == 'perf-fyi':
    builders = bot_platforms.ALL_PERF_FYI_PLATFORMS
  else:
    builders = bot_platforms.ALL_PLATFORMS
    _PromptWarning()

  if not args.use_old_timing_data:
    print 'Update shards timing data. May take a while...'
    load_timing_args = []
    for b in builders:
      load_timing_args.append((b.name, b.timing_file_path))
    p = multiprocessing.Pool(len(load_timing_args))
    p.map(_LoadTimingData, load_timing_args)

  for b in builders:
    _GenerateShardMap(
        b, b.num_shards, b.shards_map_file_path, args.debug, benchmark=None)
    print 'Updated sharding map for %s' % repr(b.name)


def _CreateShardMapForBenchmark(args):
  """Create the shard map for the given benchmark.

  Args:
    args(Namespace object): the namespace object for the subparser `create`. It
      will contain the attributes:
        `benchmark`: the name of the benchmark that we want the shard for
        `num_shards`: the total number of shards that we want to use
        `output_path`: the output file path for the shard map
        `builder`: the builder name, unlike the above, this is a string instead
          of a list of string like above
  """
  builder = None
  if args.timing_data_source:
    [builder] = [b for b in bot_platforms.ALL_PLATFORMS
                 if b.name == args.timing_data_source]
  _GenerateShardMap(
      builder, args.shards_num, args.output_path, args.debug, args.benchmark)


def main():
  parser = GetParser()
  options = parser.parse_args()
  if options.builders and options.waterfall:
    parser.error('Cannot specify both --builders and --waterfall '
                 'at the same time')

  options.func(options)

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