#!/usr/bin/env python
# Copyright 2015 gRPC authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Run test matrix."""

from __future__ import print_function

import argparse
import multiprocessing
import os
import sys

import python_utils.jobset as jobset
import python_utils.report_utils as report_utils
from python_utils.filter_pull_request_tests import filter_tests

_ROOT = os.path.abspath(os.path.join(os.path.dirname(sys.argv[0]), '../..'))
os.chdir(_ROOT)

_DEFAULT_RUNTESTS_TIMEOUT = 1*60*60

# Set the timeout high to allow enough time for sanitizers and pre-building
# clang docker.
_CPP_RUNTESTS_TIMEOUT = 4*60*60

# C++ TSAN takes longer than other sanitizers
_CPP_TSAN_RUNTESTS_TIMEOUT = 8*60*60

# Number of jobs assigned to each run_tests.py instance
_DEFAULT_INNER_JOBS = 2

# report suffix is important for reports to get picked up by internal CI
_REPORT_SUFFIX = 'sponge_log.xml'


def _report_filename(name):
  """Generates report file name"""
  return 'report_%s_%s' % (name, _REPORT_SUFFIX)


def _report_filename_internal_ci(name):
  """Generates report file name that leads to better presentation by internal CI"""
  return '%s/%s' % (name, _REPORT_SUFFIX)


def _docker_jobspec(name, runtests_args=[], runtests_envs={},
                    inner_jobs=_DEFAULT_INNER_JOBS,
                    timeout_seconds=None):
  """Run a single instance of run_tests.py in a docker container"""
  if not timeout_seconds:
    timeout_seconds = _DEFAULT_RUNTESTS_TIMEOUT
  test_job = jobset.JobSpec(
          cmdline=['python', 'tools/run_tests/run_tests.py',
                   '--use_docker',
                   '-t',
                   '-j', str(inner_jobs),
                   '-x', _report_filename(name),
                   '--report_suite_name', '%s' % name] + runtests_args,
          environ=runtests_envs,
          shortname='run_tests_%s' % name,
          timeout_seconds=timeout_seconds)
  return test_job


def _workspace_jobspec(name, runtests_args=[], workspace_name=None,
                       runtests_envs={}, inner_jobs=_DEFAULT_INNER_JOBS,
                       timeout_seconds=None):
  """Run a single instance of run_tests.py in a separate workspace"""
  if not workspace_name:
    workspace_name = 'workspace_%s' % name
  if not timeout_seconds:
    timeout_seconds = _DEFAULT_RUNTESTS_TIMEOUT
  env = {'WORKSPACE_NAME': workspace_name}
  env.update(runtests_envs)
  test_job = jobset.JobSpec(
          cmdline=['bash',
                   'tools/run_tests/helper_scripts/run_tests_in_workspace.sh',
                   '-t',
                   '-j', str(inner_jobs),
                   '-x', '../%s' % _report_filename(name),
                   '--report_suite_name', '%s' % name] + runtests_args,
          environ=env,
          shortname='run_tests_%s' % name,
          timeout_seconds=timeout_seconds)
  return test_job


def _generate_jobs(languages, configs, platforms, iomgr_platform = 'native',
                  arch=None, compiler=None,
                  labels=[], extra_args=[], extra_envs={},
                  inner_jobs=_DEFAULT_INNER_JOBS,
                  timeout_seconds=None):
  result = []
  for language in languages:
    for platform in platforms:
      for config in configs:
        name = '%s_%s_%s_%s' % (language, platform, config, iomgr_platform)
        runtests_args = ['-l', language,
                         '-c', config,
                         '--iomgr_platform', iomgr_platform]
        if arch or compiler:
          name += '_%s_%s' % (arch, compiler)
          runtests_args += ['--arch', arch,
                            '--compiler', compiler]
        if '--build_only' in extra_args:
          name += '_buildonly'
        for extra_env in extra_envs:
          name += '_%s_%s' % (extra_env, extra_envs[extra_env])

        runtests_args += extra_args
        if platform == 'linux':
          job = _docker_jobspec(name=name, runtests_args=runtests_args,
                                runtests_envs=extra_envs, inner_jobs=inner_jobs,
                                timeout_seconds=timeout_seconds)
        else:
          job = _workspace_jobspec(name=name, runtests_args=runtests_args,
                                   runtests_envs=extra_envs, inner_jobs=inner_jobs,
                                   timeout_seconds=timeout_seconds)

        job.labels = [platform, config, language, iomgr_platform] + labels
        result.append(job)
  return result


def _create_test_jobs(extra_args=[], inner_jobs=_DEFAULT_INNER_JOBS):
  test_jobs = []
  # supported on linux only
  test_jobs += _generate_jobs(languages=['sanity', 'php7'],
                             configs=['dbg', 'opt'],
                             platforms=['linux'],
                             labels=['basictests', 'multilang'],
                             extra_args=extra_args,
                             inner_jobs=inner_jobs)

  # supported on all platforms.
  test_jobs += _generate_jobs(languages=['c'],
                             configs=['dbg', 'opt'],
                             platforms=['linux', 'macos', 'windows'],
                             labels=['basictests', 'corelang'],
                             extra_args=extra_args,
                             inner_jobs=inner_jobs,
                             timeout_seconds=_CPP_RUNTESTS_TIMEOUT)
  
  test_jobs += _generate_jobs(languages=['csharp', 'node', 'python'],
                             configs=['dbg', 'opt'],
                             platforms=['linux', 'macos', 'windows'],
                             labels=['basictests', 'multilang'],
                             extra_args=extra_args,
                             inner_jobs=inner_jobs)

  # supported on linux and mac.
  test_jobs += _generate_jobs(languages=['c++'],
                              configs=['dbg', 'opt'],
                              platforms=['linux', 'macos'],
                              labels=['basictests', 'corelang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs,
                              timeout_seconds=_CPP_RUNTESTS_TIMEOUT)
  
  test_jobs += _generate_jobs(languages=['grpc-node', 'ruby', 'php'],
                              configs=['dbg', 'opt'],
                              platforms=['linux', 'macos'],
                              labels=['basictests', 'multilang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  # supported on mac only.
  test_jobs += _generate_jobs(languages=['objc'],
                              configs=['dbg', 'opt'],
                              platforms=['macos'],
                              labels=['basictests', 'multilang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  # sanitizers
  test_jobs += _generate_jobs(languages=['c'],
                              configs=['msan', 'asan', 'tsan', 'ubsan'],
                              platforms=['linux'],
                              labels=['sanitizers', 'corelang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs,
                              timeout_seconds=_CPP_RUNTESTS_TIMEOUT)
  test_jobs += _generate_jobs(languages=['c++'],
                              configs=['asan'],
                              platforms=['linux'],
                              labels=['sanitizers', 'corelang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs,
                              timeout_seconds=_CPP_RUNTESTS_TIMEOUT)
  test_jobs += _generate_jobs(languages=['c++'],
                              configs=['tsan'],
                              platforms=['linux'],
                              labels=['sanitizers', 'corelang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs,
                              timeout_seconds=_CPP_TSAN_RUNTESTS_TIMEOUT)

  return test_jobs


def _create_portability_test_jobs(extra_args=[], inner_jobs=_DEFAULT_INNER_JOBS):
  test_jobs = []
  # portability C x86
  test_jobs += _generate_jobs(languages=['c'],
                              configs=['dbg'],
                              platforms=['linux'],
                              arch='x86',
                              compiler='default',
                              labels=['portability', 'corelang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  # portability C and C++ on x64
  for compiler in ['gcc4.8', 'gcc5.3', 'gcc_musl',
                   'clang3.5', 'clang3.6', 'clang3.7']:
    test_jobs += _generate_jobs(languages=['c', 'c++'],
                                configs=['dbg'],
                                platforms=['linux'],
                                arch='x64',
                                compiler=compiler,
                                labels=['portability', 'corelang'],
                                extra_args=extra_args,
                                inner_jobs=inner_jobs,
                                timeout_seconds=_CPP_RUNTESTS_TIMEOUT)

  # portability C on Windows 64-bit (x86 is the default)
  test_jobs += _generate_jobs(languages=['c'],
                              configs=['dbg'],
                              platforms=['windows'],
                              arch='x64',
                              compiler='default',
                              labels=['portability', 'corelang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  # portability C++ on Windows
  # TODO(jtattermusch): some of the tests are failing, so we force --build_only
  test_jobs += _generate_jobs(languages=['c++'],
                              configs=['dbg'],
                              platforms=['windows'],
                              arch='default',
                              compiler='default',
                              labels=['portability', 'corelang'],
                              extra_args=extra_args + ['--build_only'],
                              inner_jobs=inner_jobs)

  # portability C and C++ on Windows using VS2017 (build only)
  # TODO(jtattermusch): some of the tests are failing, so we force --build_only
  test_jobs += _generate_jobs(languages=['c', 'c++'],
                              configs=['dbg'],
                              platforms=['windows'],
                              arch='x64',
                              compiler='cmake_vs2017',
                              labels=['portability', 'corelang'],
                              extra_args=extra_args + ['--build_only'],
                              inner_jobs=inner_jobs)

  # C and C++ with the c-ares DNS resolver on Linux
  test_jobs += _generate_jobs(languages=['c', 'c++'],
                              configs=['dbg'], platforms=['linux'],
                              labels=['portability', 'corelang'],
                              extra_args=extra_args,
                              extra_envs={'GRPC_DNS_RESOLVER': 'ares'},
                              timeout_seconds=_CPP_RUNTESTS_TIMEOUT)

  # TODO(zyc): Turn on this test after adding c-ares support on windows.
  # C with the c-ares DNS resolver on Windows
  # test_jobs += _generate_jobs(languages=['c'],
  #                             configs=['dbg'], platforms=['windows'],
  #                             labels=['portability', 'corelang'],
  #                             extra_args=extra_args,
  #                             extra_envs={'GRPC_DNS_RESOLVER': 'ares'})

  # C and C++ build with cmake on Linux
  # TODO(jtattermusch): some of the tests are failing, so we force --build_only
  # to make sure it's buildable at least.
  test_jobs += _generate_jobs(languages=['c', 'c++'],
                              configs=['dbg'],
                              platforms=['linux'],
                              arch='default',
                              compiler='cmake',
                              labels=['portability', 'corelang'],
                              extra_args=extra_args + ['--build_only'],
                              inner_jobs=inner_jobs)

  test_jobs += _generate_jobs(languages=['python'],
                              configs=['dbg'],
                              platforms=['linux'],
                              arch='default',
                              compiler='python_alpine',
                              labels=['portability', 'multilang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  test_jobs += _generate_jobs(languages=['csharp'],
                              configs=['dbg'],
                              platforms=['linux'],
                              arch='default',
                              compiler='coreclr',
                              labels=['portability', 'multilang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  test_jobs += _generate_jobs(languages=['c'],
                              configs=['dbg'],
                              platforms=['linux'],
                              iomgr_platform='uv',
                              labels=['portability', 'corelang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs,
                              timeout_seconds=_CPP_RUNTESTS_TIMEOUT)

  test_jobs += _generate_jobs(languages=['node'],
                              configs=['dbg'],
                              platforms=['linux'],
                              arch='default',
                              compiler='electron1.6',
                              labels=['portability', 'multilang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  test_jobs += _generate_jobs(languages=['node'],
                              configs=['dbg'],
                              platforms=['linux'],
                              arch='default',
                              compiler='node4',
                              labels=['portability', 'multilang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  test_jobs += _generate_jobs(languages=['node'],
                              configs=['dbg'],
                              platforms=['linux'],
                              arch='default',
                              compiler='node6',
                              labels=['portability', 'multilang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  test_jobs += _generate_jobs(languages=['node'],
                              configs=['dbg'],
                              platforms=['linux'],
                              arch='default',
                              compiler='node7',
                              labels=['portability', 'multilang'],
                              extra_args=extra_args,
                              inner_jobs=inner_jobs)

  return test_jobs


def _allowed_labels():
  """Returns a list of existing job labels."""
  all_labels = set()
  for job in _create_test_jobs() + _create_portability_test_jobs():
    for label in job.labels:
      all_labels.add(label)
  return sorted(all_labels)


def _runs_per_test_type(arg_str):
  """Auxiliary function to parse the "runs_per_test" flag."""
  try:
    n = int(arg_str)
    if n <= 0: raise ValueError
    return n
  except:
    msg = '\'{}\' is not a positive integer'.format(arg_str)
    raise argparse.ArgumentTypeError(msg)


if __name__ == "__main__":
  argp = argparse.ArgumentParser(description='Run a matrix of run_tests.py tests.')
  argp.add_argument('-j', '--jobs',
                    default=multiprocessing.cpu_count()/_DEFAULT_INNER_JOBS,
                    type=int,
                    help='Number of concurrent run_tests.py instances.')
  argp.add_argument('-f', '--filter',
                    choices=_allowed_labels(),
                    nargs='+',
                    default=[],
                    help='Filter targets to run by label with AND semantics.')
  argp.add_argument('--exclude',
                    choices=_allowed_labels(),
                    nargs='+',
                    default=[],
                    help='Exclude targets with any of given labels.')
  argp.add_argument('--build_only',
                    default=False,
                    action='store_const',
                    const=True,
                    help='Pass --build_only flag to run_tests.py instances.')
  argp.add_argument('--force_default_poller', default=False, action='store_const', const=True,
                    help='Pass --force_default_poller to run_tests.py instances.')
  argp.add_argument('--dry_run',
                    default=False,
                    action='store_const',
                    const=True,
                    help='Only print what would be run.')
  argp.add_argument('--filter_pr_tests',
                    default=False,
                    action='store_const',
                    const=True,
                    help='Filters out tests irrelevant to pull request changes.')
  argp.add_argument('--base_branch',
                    default='origin/master',
                    type=str,
                    help='Branch that pull request is requesting to merge into')
  argp.add_argument('--inner_jobs',
                    default=_DEFAULT_INNER_JOBS,
                    type=int,
                    help='Number of jobs in each run_tests.py instance')
  argp.add_argument('-n', '--runs_per_test', default=1, type=_runs_per_test_type,
                    help='How many times to run each tests. >1 runs implies ' +
                    'omitting passing test from the output & reports.')
  argp.add_argument('--max_time', default=-1, type=int,
                    help='Maximum amount of time to run tests for' +
                         '(other tests will be skipped)')
  argp.add_argument('--internal_ci',
                    default=False,
                    action='store_const',
                    const=True,
                    help='Put reports into subdirectories to improve presentation of '
                    'results by Internal CI.')
  argp.add_argument('--bq_result_table',
                    default='',
                    type=str,
                    nargs='?',
                    help='Upload test results to a specified BQ table.')
  args = argp.parse_args()

  if args.internal_ci:
    _report_filename = _report_filename_internal_ci  # override the function

  extra_args = []
  if args.build_only:
    extra_args.append('--build_only')
  if args.force_default_poller:
    extra_args.append('--force_default_poller')
  if args.runs_per_test > 1:
    extra_args.append('-n')
    extra_args.append('%s' % args.runs_per_test)
    extra_args.append('--quiet_success')
  if args.max_time > 0:
    extra_args.extend(('--max_time', '%d' % args.max_time))
  if args.bq_result_table:
    extra_args.append('--bq_result_table')
    extra_args.append('%s' % args.bq_result_table)
    extra_args.append('--measure_cpu_costs')
    extra_args.append('--disable_auto_set_flakes')

  all_jobs = _create_test_jobs(extra_args=extra_args, inner_jobs=args.inner_jobs) + \
             _create_portability_test_jobs(extra_args=extra_args, inner_jobs=args.inner_jobs)

  jobs = []
  for job in all_jobs:
    if not args.filter or all(filter in job.labels for filter in args.filter):
      if not any(exclude_label in job.labels for exclude_label in args.exclude):
        jobs.append(job)

  if not jobs:
    jobset.message('FAILED', 'No test suites match given criteria.',
                   do_newline=True)
    sys.exit(1)

  print('IMPORTANT: The changes you are testing need to be locally committed')
  print('because only the committed changes in the current branch will be')
  print('copied to the docker environment or into subworkspaces.')

  skipped_jobs = []

  if args.filter_pr_tests:
    print('Looking for irrelevant tests to skip...')
    relevant_jobs = filter_tests(jobs, args.base_branch)
    if len(relevant_jobs) == len(jobs):
      print('No tests will be skipped.')
    else:
      print('These tests will be skipped:')
      skipped_jobs = list(set(jobs) - set(relevant_jobs))
      # Sort by shortnames to make printing of skipped tests consistent
      skipped_jobs.sort(key=lambda job: job.shortname)
      for job in list(skipped_jobs):
        print('  %s' % job.shortname)
    jobs = relevant_jobs

  print('Will run these tests:')
  for job in jobs:
    if args.dry_run:
      print('  %s: "%s"' % (job.shortname, ' '.join(job.cmdline)))
    else:
      print('  %s' % job.shortname)
  print

  if args.dry_run:
    print('--dry_run was used, exiting')
    sys.exit(1)

  jobset.message('START', 'Running test matrix.', do_newline=True)
  num_failures, resultset = jobset.run(jobs,
                                       newline_on_success=True,
                                       travis=True,
                                       maxjobs=args.jobs)
  # Merge skipped tests into results to show skipped tests on report.xml
  if skipped_jobs:
    ignored_num_skipped_failures, skipped_results = jobset.run(
        skipped_jobs, skip_jobs=True)
    resultset.update(skipped_results)
  report_utils.render_junit_xml_report(resultset, _report_filename('aggregate_tests'),
                                       suite_name='aggregate_tests')

  if num_failures == 0:
    jobset.message('SUCCESS', 'All run_tests.py instance finished successfully.',
                   do_newline=True)
  else:
    jobset.message('FAILED', 'Some run_tests.py instance have failed.',
                   do_newline=True)
    sys.exit(1)