# -*- coding: utf-8 -*-

import sys
import gevent
import json
import signal
import logging
import time
from locust.env import Environment
from locust.log import setup_logging
from locust.stats import stats_printer
from locust.util.timespan import parse_timespan

setup_logging("INFO", None)
logger = logging.getLogger(__name__)


def sig_term_handler():
    logger.info("Got SIGTERM signal")
    sys.exit(0)


class LocustLoadTest(object):
    """
    Runs a Locust load test and returns statistics
    """

    def __init__(self, settings):
        self.settings = settings
        self.start_time = None
        self.end_time = None
        gevent.signal_handler(signal.SIGTERM, sig_term_handler)

    def stats(self):
        """
        Returns the statistics from the load test in JSON
        """
        statistics = {
            "requests": {},
            "failures": {},
            "num_requests": self.env.runner.stats.num_requests,
            "num_requests_fail": self.env.runner.stats.num_failures,
            "start_time": self.start_time,
            "end_time": self.end_time,
        }

        for name, value in self.env.runner.stats.entries.items():
            locust_task_name = "{0}_{1}".format(name[1], name[0])
            statistics["requests"][locust_task_name] = {
                "request_type": name[1],
                "num_requests": value.num_requests,
                "min_response_time": value.min_response_time,
                "median_response_time": value.median_response_time,
                "avg_response_time": value.avg_response_time,
                "max_response_time": value.max_response_time,
                "response_times": value.response_times,
                "response_time_percentiles": {
                    55: value.get_response_time_percentile(0.55),
                    65: value.get_response_time_percentile(0.65),
                    75: value.get_response_time_percentile(0.75),
                    85: value.get_response_time_percentile(0.85),
                    95: value.get_response_time_percentile(0.95),
                },
                "total_rps": value.total_rps,
                "total_rpm": value.total_rps * 60,
            }

        for id, error in self.env.runner.errors.items():
            error_dict = error.to_dict()
            locust_task_name = "{0}_{1}".format(
                error_dict["method"], error_dict["name"]
            )
            statistics["failures"][locust_task_name] = error_dict

        return statistics

    def set_run_time_in_sec(self, run_time_str):
        try:
            self.run_time_in_sec = parse_timespan(run_time_str)
        except ValueError:
            logger.error(
                "Invalid format for `run_time` parameter: '%s', "
                "Valid formats are: 20s, 3m, 2h, 1h20m, 3h30m10s, etc." % run_time_str
            )
            sys.exit(1)
        except TypeError:
            logger.error(
                "`run_time` must be a string, not %s. Received value: % "
                % (type(run_time_str), run_time_str)
            )
            sys.exit(1)

    def run(self):
        """
        Run the load test.
        """

        if self.settings.run_time:
            self.set_run_time_in_sec(run_time_str=self.settings.run_time)

            logger.info("Run time limit set to %s seconds" % self.run_time_in_sec)

            def timelimit_stop():
                logger.info(
                    "Run time limit reached: %s seconds. Stopping Locust Runner."
                    % self.run_time_in_sec
                )
                self.env.runner.quit()
                self.end_time = time.time()
                logger.info(
                    "Locust completed %s requests with %s errors"
                    % (self.env.runner.stats.num_requests, len(self.env.runner.errors))
                )
                logger.info(json.dumps(self.stats()))

            gevent.spawn_later(self.run_time_in_sec, timelimit_stop)

        try:
            logger.info("Starting Locust with settings %s " % vars(self.settings))

            self.env = Environment(
                user_classes=self.settings.classes,
                host=self.settings.host,
                tags=self.settings.tags,
                exclude_tags=self.settings.exclude_tags,
                reset_stats=self.settings.reset_stats,
                step_load=self.settings.step_load,
                stop_timeout=self.settings.stop_timeout,
            )

            self.env.create_local_runner()
            gevent.spawn(stats_printer(self.env.stats))

            self.env.runner.start(
                user_count=self.settings.num_users, hatch_rate=self.settings.hatch_rate
            )

            self.start_time = time.time()
            self.env.runner.greenlet.join()

        except Exception as e:
            logger.error("Locust exception {0}".format(repr(e)))

        finally:
            self.env.events.quitting.fire()