# Copyright 2016 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.

# performance scenario configuration for various languages

import math

WARMUP_SECONDS = 5
JAVA_WARMUP_SECONDS = 15  # Java needs more warmup time for JIT to kick in.
BENCHMARK_SECONDS = 30

SMOKETEST = 'smoketest'
SCALABLE = 'scalable'
INPROC = 'inproc'
SWEEP = 'sweep'
DEFAULT_CATEGORIES = (SCALABLE, SMOKETEST)

SECURE_SECARGS = {
    'use_test_ca': True,
    'server_host_override': 'foo.test.google.fr'
}

HISTOGRAM_PARAMS = {
    'resolution': 0.01,
    'max_possible': 60e9,
}

# target number of RPCs outstanding on across all client channels in
# non-ping-pong tests (since we can only specify per-channel numbers, the
# actual target will be slightly higher)
OUTSTANDING_REQUESTS = {'async': 6400, 'async-limited': 800, 'sync': 1000}

# wide is the number of client channels in multi-channel tests (1 otherwise)
WIDE = 64


def _get_secargs(is_secure):
    if is_secure:
        return SECURE_SECARGS
    else:
        return None


def remove_nonproto_fields(scenario):
    """Removes special-purpose fields that don't belong in the protobuf.

    This function removes additional information about the scenario that is not
    included in the ScenarioConfig protobuf message.
    """
    scenario.pop('CATEGORIES', None)
    scenario.pop('CLIENT_LANGUAGE', None)
    scenario.pop('SERVER_LANGUAGE', None)
    scenario.pop('EXCLUDED_POLL_ENGINES', None)
    return scenario


def geometric_progression(start, stop, step):
    n = start
    while n < stop:
        yield int(round(n))
        n *= step


def _payload_type(use_generic_payload, req_size, resp_size):
    r = {}
    sizes = {
        'req_size': req_size,
        'resp_size': resp_size,
    }
    if use_generic_payload:
        r['bytebuf_params'] = sizes
    else:
        r['simple_params'] = sizes
    return r


def _load_params(offered_load):
    r = {}
    if offered_load is None:
        r['closed_loop'] = {}
    else:
        load = {}
        load['offered_load'] = offered_load
        r['poisson'] = load
    return r


def _add_channel_arg(config, key, value):
    if 'channel_args' in config:
        channel_args = config['channel_args']
    else:
        channel_args = []
        config['channel_args'] = channel_args
    arg = {'name': key}
    if isinstance(value, int):
        arg['int_value'] = value
    else:
        arg['str_value'] = value
    channel_args.append(arg)


def _ping_pong_scenario(name,
                        rpc_type,
                        client_type,
                        server_type,
                        secure=True,
                        use_generic_payload=False,
                        req_size=0,
                        resp_size=0,
                        unconstrained_client=None,
                        client_language=None,
                        server_language=None,
                        async_server_threads=0,
                        client_processes=0,
                        server_processes=0,
                        server_threads_per_cq=0,
                        client_threads_per_cq=0,
                        warmup_seconds=WARMUP_SECONDS,
                        categories=None,
                        channels=None,
                        outstanding=None,
                        num_clients=None,
                        resource_quota_size=None,
                        messages_per_stream=None,
                        excluded_poll_engines=None,
                        minimal_stack=False,
                        offered_load=None):
    """Creates a basic ping pong scenario."""
    scenario = {
        'name': name,
        'num_servers': 1,
        'num_clients': 1,
        'client_config': {
            'client_type': client_type,
            'security_params': _get_secargs(secure),
            'outstanding_rpcs_per_channel': 1,
            'client_channels': 1,
            'async_client_threads': 1,
            'client_processes': client_processes,
            'threads_per_cq': client_threads_per_cq,
            'rpc_type': rpc_type,
            'histogram_params': HISTOGRAM_PARAMS,
            'channel_args': [],
        },
        'server_config': {
            'server_type': server_type,
            'security_params': _get_secargs(secure),
            'async_server_threads': async_server_threads,
            'server_processes': server_processes,
            'threads_per_cq': server_threads_per_cq,
            'channel_args': [],
        },
        'warmup_seconds': warmup_seconds,
        'benchmark_seconds': BENCHMARK_SECONDS,
        'CATEGORIES': list(DEFAULT_CATEGORIES),
        'EXCLUDED_POLL_ENGINES': [],
    }
    if resource_quota_size:
        scenario['server_config']['resource_quota_size'] = resource_quota_size
    if use_generic_payload:
        if server_type != 'ASYNC_GENERIC_SERVER':
            raise Exception('Use ASYNC_GENERIC_SERVER for generic payload.')
        scenario['server_config']['payload_config'] = _payload_type(
            use_generic_payload, req_size, resp_size)

    scenario['client_config']['payload_config'] = _payload_type(
        use_generic_payload, req_size, resp_size)

    # Optimization target of 'throughput' does not work well with epoll1 polling
    # engine. Use the default value of 'blend'
    optimization_target = 'throughput'

    if unconstrained_client:
        outstanding_calls = outstanding if outstanding is not None else OUTSTANDING_REQUESTS[
            unconstrained_client]
        # clamp buffer usage to something reasonable (16 gig for now)
        MAX_MEMORY_USE = 16 * 1024 * 1024 * 1024
        if outstanding_calls * max(req_size, resp_size) > MAX_MEMORY_USE:
            outstanding_calls = max(1,
                                    MAX_MEMORY_USE / max(req_size, resp_size))
        wide = channels if channels is not None else WIDE
        deep = int(math.ceil(1.0 * outstanding_calls / wide))

        scenario[
            'num_clients'] = num_clients if num_clients is not None else 0  # use as many clients as available.
        scenario['client_config']['outstanding_rpcs_per_channel'] = deep
        scenario['client_config']['client_channels'] = wide
        scenario['client_config']['async_client_threads'] = 0
        if offered_load is not None:
            optimization_target = 'latency'
    else:
        scenario['client_config']['outstanding_rpcs_per_channel'] = 1
        scenario['client_config']['client_channels'] = 1
        scenario['client_config']['async_client_threads'] = 1
        optimization_target = 'latency'

    scenario['client_config']['load_params'] = _load_params(offered_load)

    optimization_channel_arg = {
        'name': 'grpc.optimization_target',
        'str_value': optimization_target
    }
    scenario['client_config']['channel_args'].append(optimization_channel_arg)
    scenario['server_config']['channel_args'].append(optimization_channel_arg)

    if minimal_stack:
        _add_channel_arg(scenario['client_config'], 'grpc.minimal_stack', 1)
        _add_channel_arg(scenario['server_config'], 'grpc.minimal_stack', 1)

    if messages_per_stream: