# Copyright 2014: Mirantis Inc.
# All Rights Reserved.
#
#    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.

import collections
import multiprocessing
import threading
import time

from rally.common import log as logging
from rally.common import utils
from rally import consts
from rally.task import runner

LOG = logging.getLogger(__name__)


def _worker_process(queue, iteration_gen, timeout, rps, times,
                    max_concurrent, context, cls, method_name,
                    args, aborted, info):
    """Start scenario within threads.

    Spawn N threads per second. Each thread runs the scenario once, and appends
    result to queue. A maximum of max_concurrent threads will be ran
    concurrently.

    :param queue: queue object to append results
    :param iteration_gen: next iteration number generator
    :param timeout: operation's timeout
    :param rps: number of scenario iterations to be run per one second
    :param times: total number of scenario iterations to be run
    :param max_concurrent: maximum worker concurrency
    :param context: scenario context object
    :param cls: scenario class
    :param method_name: scenario method name
    :param args: scenario args
    :param aborted: multiprocessing.Event that aborts load generation if
                    the flag is set
    :param info: info about all processes count and counter of runned process
    """

    pool = collections.deque()
    sleep = 1.0 / rps

    runner._log_worker_info(times=times, rps=rps, timeout=timeout,
                            cls=cls, method_name=method_name, args=args)

    time.sleep(
        (sleep * info["processes_counter"]) / info["processes_to_start"])

    start = time.time()

    i = 0
    while i < times and not aborted.is_set():
        scenario_context = runner._get_scenario_context(context)
        scenario_args = (next(iteration_gen), cls, method_name,
                         scenario_context, args)
        worker_args = (queue, scenario_args)
        thread = threading.Thread(target=runner._worker_thread,
                                  args=worker_args)
        i += 1
        thread.start()
        pool.append(thread)

        time_gap = time.time() - start
        real_rps = i / time_gap if time_gap else "Infinity"

        LOG.debug("Worker: %s rps: %s (requested rps: %s)" %
                  (i, real_rps, rps))

        # try to join latest thread(s) until it finished, or until time to
        # start new thread (if we have concurrent slots available)
        while i / (time.time() - start) > rps or len(pool) >= max_concurrent:
            if pool:
                pool[0].join(0.001)
                if not pool[0].isAlive():
                    pool.popleft()
            else:
                time.sleep(0.001)

    while pool:
        thr = pool.popleft()
        thr.join()


@runner.configure(name="rps")
class RPSScenarioRunner(runner.ScenarioRunner):
    """Scenario runner that does the job with specified frequency.

    Every single benchmark scenario iteration is executed with specified
    frequency (runs per second) in a pool of processes. The scenario will be
    launched for a fixed number of times in total (specified in the config).

    An example of a rps scenario is booting 1 VM per second. This
    execution type is thus very helpful in understanding the maximal load that
    a certain cloud can handle.
    """

    CONFIG_SCHEMA = {
        "type": "object",
        "$schema": consts.JSON_SCHEMA,
        "properties": {
            "type": {
                "type": "string"
            },
            "times": {
                "type": "integer",
                "minimum": 1
            },
            "rps": {
                "type": "number",
                "exclusiveMinimum": True,
                "minimum": 0
            },
            "timeout": {
                "type": "number",
            },
            "max_concurrency": {
                "type": "integer",
                "minimum": 1
            },
            "max_cpu_count": {
                "type": "integer",
                "minimum": 1
            }
        },
        "additionalProperties": False
    }

    def _run_scenario(self, cls, method_name, context, args):
        """Runs the specified benchmark scenario with given arguments.

        Every single benchmark scenario iteration is executed with specified
        frequency (runs per second) in a pool of processes. The scenario will
        be launched for a fixed number of times in total (specified in the
        config).

        :param cls: The Scenario class where the scenario is implemented
        :param method_name: Name of the method that implements the scenario
        :param context: Benchmark context that contains users, admin & other
                        information, that was created before benchmark started.
        :param args: Arguments to call the scenario method with

        :returns: List of results fore each single scenario iteration,
                  where each result is a dictionary
        """
        times = self.config["times"]
        timeout = self.config.get("timeout", 0)  # 0 means no timeout
        iteration_gen = utils.RAMInt()

        cpu_count = multiprocessing.cpu_count()
        max_cpu_used = min(cpu_count,
                           self.config.get("max_cpu_count", cpu_count))

        processes_to_start = min(max_cpu_used, times,
                                 self.config.get("max_concurrency", times))
        rps_per_worker = float(self.config["rps"]) / processes_to_start
        times_per_worker, times_overhead = divmod(times, processes_to_start)

        # Determine concurrency per worker
        concurrency_per_worker, concurrency_overhead = divmod(
            self.config.get("max_concurrency", times), processes_to_start)

        self._log_debug_info(times=times, timeout=timeout,
                             max_cpu_used=max_cpu_used,
                             processes_to_start=processes_to_start,
                             rps_per_worker=rps_per_worker,
                             times_per_worker=times_per_worker,
                             times_overhead=times_overhead,
                             concurrency_per_worker=concurrency_per_worker,
                             concurrency_overhead=concurrency_overhead)

        result_queue = multiprocessing.Queue()

        def worker_args_gen(times_overhead, concurrency_overhead):
            """Generate arguments for process worker.

            Remainder of threads per process division is distributed to
            process workers equally - one thread per each process worker
            until the remainder equals zero. The same logic is applied
            to concurrency overhead.
            :param times_overhead: remaining number of threads to be
                                   distributed to workers
            :param concurrency_overhead: remaining number of maximum
                                         concurrent threads to be distributed
                                         to workers
            """
            while True:
                yield (result_queue, iteration_gen, timeout, rps_per_worker,
                       times_per_worker + (times_overhead and 1),
                       concurrency_per_worker + (concurrency_overhead and 1),
                       context, cls, method_name, args, self.aborted)
                if times_overhead:
                    times_overhead -= 1
                if concurrency_overhead:
                    concurrency_overhead -= 1

        process_pool = self._create_process_pool(
            processes_to_start, _worker_process,
            worker_args_gen(times_overhead, concurrency_overhead))
        self._join_processes(process_pool, result_queue)
