#!/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 tests in parallel."""

from __future__ import print_function

import argparse
import ast
import collections
import glob
import itertools
import json
import logging
import multiprocessing
import os
import os.path
import pipes
import platform
import random
import re
import socket
import subprocess
import sys
import tempfile
import traceback
import time
from six.moves import urllib
import uuid
import six

import python_utils.jobset as jobset
import python_utils.report_utils as report_utils
import python_utils.watch_dirs as watch_dirs
import python_utils.start_port_server as start_port_server
try:
    from python_utils.upload_test_results import upload_results_to_bq
except (ImportError):
    pass  # It's ok to not import because this is only necessary to upload results to BQ.

gcp_utils_dir = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '../gcp/utils'))
sys.path.append(gcp_utils_dir)

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

_FORCE_ENVIRON_FOR_WRAPPERS = {
    'GRPC_VERBOSITY': 'DEBUG',
}

_POLLING_STRATEGIES = {
    'linux': ['epollex', 'epoll1', 'poll'],
    'mac': ['poll'],
}

BigQueryTestData = collections.namedtuple('BigQueryTestData', 'name flaky cpu')


def get_bqtest_data(limit=None):
    import big_query_utils

    bq = big_query_utils.create_big_query()
    query = """
SELECT
  filtered_test_name,
  SUM(result != 'PASSED' AND result != 'SKIPPED') > 0 as flaky,
  MAX(cpu_measured) + 0.01 as cpu
  FROM (
  SELECT
    REGEXP_REPLACE(test_name, r'/\d+', '') AS filtered_test_name,
    result, cpu_measured
  FROM
    [grpc-testing:jenkins_test_results.aggregate_results]
  WHERE
    timestamp >= DATE_ADD(CURRENT_DATE(), -1, "WEEK")
    AND platform = '""" + platform_string() + """'
    AND NOT REGEXP_MATCH(job_name, '.*portability.*') )
GROUP BY
  filtered_test_name"""
    if limit:
        query += " limit {}".format(limit)
    query_job = big_query_utils.sync_query_job(bq, 'grpc-testing', query)
    page = bq.jobs().getQueryResults(
        pageToken=None, **query_job['jobReference']).execute(num_retries=3)
    test_data = [
        BigQueryTestData(row['f'][0]['v'], row['f'][1]['v'] == 'true',
                         float(row['f'][2]['v'])) for row in page['rows']
    ]
    return test_data


def platform_string():
    return jobset.platform_string()


_DEFAULT_TIMEOUT_SECONDS = 5 * 60
_PRE_BUILD_STEP_TIMEOUT_SECONDS = 10 * 60


def run_shell_command(cmd, env=None, cwd=None):
    try:
        subprocess.check_output(cmd, shell=True, env=env, cwd=cwd)
    except subprocess.CalledProcessError as e:
        logging.exception(
            "Error while running command '%s'. Exit status %d. Output:\n%s",
            e.cmd, e.returncode, e.output)
        raise


def max_parallel_tests_for_current_platform():
    # Too much test parallelization has only been seen to be a problem
    # so far on windows.
    if jobset.platform_string() == 'windows':
        return 64
    return 1024


# SimpleConfig: just compile with CONFIG=config, and run the binary to test
class Config(object):

    def __init__(self,
                 config,
                 environ=None,
                 timeout_multiplier=1,
                 tool_prefix=[],
                 iomgr_platform='native'):
        if environ is None:
            environ = {}
        self.build_config = config
        self.environ = environ
        self.environ['CONFIG'] = config
        self.tool_prefix = tool_prefix
        self.timeout_multiplier = timeout_multiplier
        self.iomgr_platform = iomgr_platform

    def job_spec(self,
                 cmdline,
                 timeout_seconds=_DEFAULT_TIMEOUT_SECONDS,
                 shortname=None,
                 environ={},
                 cpu_cost=1.0,
                 flaky=False):
        """Construct a jobset.JobSpec for a test under this config

       Args:
         cmdline:      a list of strings specifying the command line the test
                       would like to run
    """
        actual_environ = self.environ.copy()
        for k, v in environ.items():
            actual_environ[k] = v
        if not flaky and shortname and shortname in flaky_tests:
            flaky = True
        if shortname in shortname_to_cpu:
            cpu_cost = shortname_to_cpu[shortname]
        return jobset.JobSpec(
            cmdline=self.tool_prefix + cmdline,
            shortname=shortname,
            environ=actual_environ,
            cpu_cost=cpu_cost,
            timeout_seconds=(self.timeout_multiplier * timeout_seconds
                             if timeout_seconds else None),
            flake_retries=4 if flaky or args.allow_flakes else 0,
            timeout_retries=1 if flaky or args.allow_flakes else 0)


def get_c_tests(travis, test_lang):
    out = []
    platforms_str = 'ci_platforms' if travis else 'platforms'
    with open('tools/run_tests/generated/tests.json') as f:
        js = json.load(f)
        return [
            tgt for tgt in js
            if tgt['language'] == test_lang and platform_string() in
            tgt[platforms_str] and not (travis and tgt['flaky'])
        ]


def _check_compiler(compiler, supported_compilers):
    if compiler not in supported_compilers:
        raise Exception(
            'Compiler %s not supported (on this platform).' % compiler)


def _check_arch(arch, supported_archs):
    if arch not in supported_archs:
        raise Exception('Architecture %s not supported.' % arch)


def _is_use_docker_child():
    """Returns True if running running as a --use_docker child."""
    return True if os.getenv('RUN_TESTS_COMMAND') else False


_PythonConfigVars = collections.namedtuple('_ConfigVars', [
    'shell',
    'builder',
    'builder_prefix_arguments',
    'venv_relative_python',
    'toolchain',
    'runner',
    'test_name',
    'iomgr_platform',
])


def _python_config_generator(name, major, minor, bits, config_vars):
    name += '_' + config_vars.iomgr_platform
    return PythonConfig(
        name, config_vars.shell + config_vars.builder +
        config_vars.builder_prefix_arguments + [
            _python_pattern_function(major=major, minor=minor, bits=bits)
        ] + [name] + config_vars.venv_relative_python + config_vars.toolchain,
        config_vars.shell + config_vars.runner + [
            os.path.join(name, config_vars.venv_relative_python[0]),
            config_vars.test_name
        ])


def _pypy_config_generator(name, major, config_vars):
    return PythonConfig(
        name,
        config_vars.shell + config_vars.builder +
        config_vars.builder_prefix_arguments + [
            _pypy_pattern_function(major=major)
        ] + [name] + config_vars.venv_relative_python + config_vars.toolchain,
        config_vars.shell + config_vars.runner +
        [os.path.join(name, config_vars.venv_relative_python[0])])


def _python_pattern_function(major, minor, bits):
    # Bit-ness is handled by the test machine's environment
    if os.name == "nt":
        if bits == "64":
            return '/c/Python{major}{minor}/python.exe'.format(
                major=major, minor=minor, bits=bits)
        else:
            return '/c/Python{major}{minor}_{bits}bits/python.exe'.format(
                major=major, minor=minor, bits=bits)
    else:
        return 'python{major}.{minor}'.format(major=major, minor=minor)


def _pypy_pattern_function(major):
    if major == '2':
        return 'pypy'
    elif major == '3':
        return 'pypy3'
    else:
        raise ValueError("Unknown PyPy major version")


class CLanguage(object):

    def __init__(self, make_target, test_lang):
        self.make_target = make_target
        self.platform = platform_string()
        self.test_lang = test_lang

    def configure(self, config, args):
        self.config = config
        self.args = args
        if self.platform == 'windows':
            _check_compiler(
                self.args.compiler,
                ['default', 'cmake', 'cmake_vs2015', 'cmake_vs2017'])
            _check_arch(self.args.arch, ['default', 'x64', 'x86'])
            self._cmake_generator_option = 'Visual Studio 15 2017' if self.args.compiler == 'cmake_vs2017' else 'Visual Studio 14 2015'
            self._cmake_arch_option = 'x64' if self.args.arch == 'x64' else 'Win32'
            self._use_cmake = True
            self._make_options = []
        elif self.args.compiler == 'cmake':
            _check_arch(self.args.arch, ['default'])
            self._use_cmake = True
            self._docker_distro = 'jessie'
            self._make_options = []
        else:
            self._use_cmake = False
            self._docker_distro, self._make_options = self._compiler_options(
                self.args.use_docker, self.args.compiler)
        if args.iomgr_platform == "uv":
            cflags = '-DGRPC_UV -DGRPC_CUSTOM_IOMGR_THREAD_CHECK -DGRPC_CUSTOM_SOCKET '
            try:
                cflags += subprocess.check_output(
                    ['pkg-config', '--cflags', 'libuv']).strip() + ' '
            except (subprocess.CalledProcessError, OSError):
                pass
            try:
                ldflags = subprocess.check_output(