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

from __future__ import annotations

import collections
import multiprocessing
import queue as Queue
import threading
import time
import typing as t

from rally.common import logging
from rally.common import utils
from rally.common import validation
from rally import consts
from rally.task import runner

if t.TYPE_CHECKING:  # pragma: no cover
    from rally.task import scenario

LOG = logging.getLogger(__name__)


def _worker_process(
    queue: multiprocessing.Queue[runner.ScenarioRunnerResult],
    iteration_gen: t.Iterator[int],
    timeout: float | None,
    times: int,
    max_concurrent: int,
    context: dict[str, t.Any],
    cls: type[scenario.Scenario],
    method_name: str,
    args: dict[str, t.Any],
    event_queue: multiprocessing.Queue[dict[str, t.Any]],
    aborted: multiprocessing.synchronize.Event,
    runs_per_second: t.Callable[[dict[str, float] | float, float, int], float],
    rps_cfg: dict[str, float] | float,
    processes_to_start: int,
    info: dict[str, t.Any]
) -> None:
    """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 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 runs_per_second: function that should return desired rps value
    :param rps_cfg: rps section from task config
    :param processes_to_start: int, number of started processes for scenario
                               execution
    :param info: info about all processes count and counter of runned process
    """

    pool: collections.deque[threading.Thread] = collections.deque()
    if isinstance(rps_cfg, dict):
        rps = rps_cfg["start"]
    else:
        rps = rps_cfg
    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()
    timeout_queue: Queue.Queue[
        tuple[threading.Thread, float] | tuple[None, None]
    ] = Queue.Queue()

    if timeout:
        collector_thr_by_timeout = threading.Thread(
            target=utils.timeout_thread,
            args=(timeout_queue, )
        )
        collector_thr_by_timeout.start()

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

        i += 1
        thread.start()
        if timeout:
            timeout_queue.put((thread, time.time() + timeout))
        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, runs_per_second(rps_cfg, start, processes_to_start)))

        # 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) > runs_per_second(
                rps_cfg, start, processes_to_start) or (
                    len(pool) >= max_concurrent):
            if pool:
                pool[0].join(0.001)
                if not pool[0].is_alive():
                    pool.popleft()
            else:
                time.sleep(0.001)

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

    if timeout:
        timeout_queue.put((None, None,))
        collector_thr_by_timeout.join()


@validation.configure("check_rps")
class CheckPRSValidator(validation.Validator):
    """Additional schema validation for rps runner"""

    def validate(
        self,
        context: dict[str, t.Any],
        config: dict[str, t.Any] | None,
        plugin_cls: type[runner.plugin.Plugin],
        plugin_cfg: dict[str, t.Any] | None
    ) -> None:
        if plugin_cfg and isinstance(plugin_cfg["rps"], dict):
            if plugin_cfg["rps"]["end"] < plugin_cfg["rps"]["start"]:
                msg = "rps end value must not be less than rps start value."
                self.fail(msg)


def _runs_per_second(
    rps_cfg: dict[str, float] | float,
    start_timer: float,
    number_of_processes: int
) -> float:
    """At the given second return desired rps."""

    if not isinstance(rps_cfg, dict):
        return float(rps_cfg) / number_of_processes
    stage_order = (time.time() - start_timer) / rps_cfg.get(
        "duration", 1) - 1
    rps = (float(rps_cfg["start"] + rps_cfg["step"] * stage_order)
           / number_of_processes)

    return min(rps, float(rps_cfg["end"]))


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

    Every single 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_SCHEMA7,
        "properties": {
            "times": {
                "type": "integer",
                "minimum": 1
            },
            "rps": {
                "anyOf": [
                    {
                        "description": "Generate constant requests per second "
                                       "during the whole workload.",
                        "type": "number",
                        "exclusiveMinimum": 0,
                        "minimum": 0
                    },
                    {
                        "type": "object",
                        "description": "Increase requests per second for "
                                       "specified value each time after a "
                                       "certain number of seconds.",
                        "properties": {
                            "start": {
                                "type": "number",
                                "minimum": 1
                            },
                            "end": {
                                "type": "number",
                                "minimum": 1
                            },
                            "step": {
                                "type": "number",
                                "minimum": 1
                            },
                            "duration": {
                                "type": "number",
                                "minimum": 1
                            }
                        },
                        "additionalProperties": False,
                        "required": ["start", "end", "step"]
                    }
                ],
            },
            "timeout": {
                "type": "number",
            },
            "max_concurrency": {
                "type": "integer",
                "minimum": 1
            },
            "max_cpu_count": {
                "type": "integer",
                "minimum": 1
            }
        },
        "required": ["times", "rps"],
        "additionalProperties": False
    }

    def _run_scenario(
        self,
        cls: type[scenario.Scenario],
        method_name: t.Literal["run"],
        context: dict[str, t.Any],
        args: dict[str, t.Any]
    ) -> None:
        """Runs the specified scenario with given arguments.

        Every single scenario iteration is executed with specified
        frequency (runs per second) in a pool of processes. The scenario is
        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: Context that contains users, admin & other
                        information, that was created before scenario
                        execution starts.
        :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))
        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,
                             times_per_worker=times_per_worker,
                             times_overhead=times_overhead,
                             concurrency_per_worker=concurrency_per_worker,
                             concurrency_overhead=concurrency_overhead)

        result_queue: multiprocessing.Queue[runner.ScenarioRunnerResult] = (
            multiprocessing.Queue())
        event_queue: multiprocessing.Queue[dict[str, t.Any]] = (
            multiprocessing.Queue())

        def worker_args_gen(
            times_overhead: int, concurrency_overhead: int
        ) -> t.Generator[tuple[t.Any, ...], None, None]:
            """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,
                    times_per_worker + (times_overhead and 1),
                    concurrency_per_worker + (concurrency_overhead and 1),
                    context, cls, method_name, args, event_queue,
                    self.aborted, _runs_per_second, self.config["rps"],
                    processes_to_start
                )
                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, event_queue)
