#!/usr/bin/python3

import argparse
import os
import subprocess
import threading
import io
import sys
import re
import logging
import time
import multiprocessing
import histogram

generate_concord_keys = "@concord_bft_tools_BINARY_DIR@/GenerateConcordKeys"
bin_dir = "@CMAKE_BINARY_DIR@/tests/simpleTest"
log_dir = "@CMAKE_BINARY_DIR@/tests/simpleTest/scripts"

g_logger = logging.getLogger("simpletest")
g_log_formatter = logging.Formatter(
    "%(asctime)-15s %(levelname)-6s %(""message)s")
g_log_console_handler = logging.StreamHandler()
g_log_console_handler.setFormatter(g_log_formatter)
g_logger.addHandler(g_log_console_handler)

class CLI:
    def __init__(self, num_it, vc_t, measure_perf):
        self.num_of_iterations = num_it
        self.vc_timeout = vc_t
        self.measure_perf = measure_perf

class SimpleClientParams:
    def __init__(self):
        self.initial_retry_timeout_milli = 1100
        self.min_retry_timeout_milli = 1000
        self.max_retry_timeout_milli = 2000
        self.sends_request_to_all_replicas_first_thresh = 2
        self.sends_request_to_all_replicas_period_thresh = 2
        self.periodic_reset_thresh = 30

class BFT:
    def __init__(self, name, n, r, f, c, cl, vc):
        self.name = name
        self.num_of_replicas = n
        self.num_of_replicas_to_run = r
        self.num_of_faulties = f
        self.num_of_slow = c
        self.num_of_clients = cl
        self.test_vc = vc
        self.vc_will_fail = False if not vc else n - r >= f


class Environment:
    def __init__(self, rep_exe, cl_exe, binary_dir, log_dir, no_console):
        self.replica_exec = rep_exe
        self.client_exec = cl_exe
        self.binary_dir = binary_dir
        self.log_dir = log_dir
        self.no_console = no_console


class TestResult:
    def __init__(self, exit_code, user_msg, summary):
        self.exit_code = exit_code
        self.user_msg = user_msg
        self.perf_summary = summary

class ProcessOutput:
    def __init__(self, exec_name, exit_code, user_msg, uid, latencies):
        self.uid = uid
        self.user_msg = user_msg
        self.exec_name = exec_name
        self.exit_code = exit_code
        self.latencies = latencies


class TestProcess:
    def __init__(self,
                 friendly_name,
                 executable,
                 params,
                 log_dir,
                 is_waitable,
                 uid,
                 no_console,
                 perf_enabled,
                 on_new_line_callback = None):
        self.friendly_name = friendly_name
        # full path to the executable
        self.executable = executable
        self.log_dir = log_dir
        self.params = params
        self.is_waitable = is_waitable
        self.uid = uid
        self.process = None
        self.thread = None
        self.out_file = None
        self.no_console = no_console
        self.perf_enabled = perf_enabled,
        self.on_new_line_callback = on_new_line_callback
        self.perf_file = None
        self.latencies = None

    def invoke_on_new_line(self, txt):
        if self.on_new_line_callback:
            self.on_new_line_callback=self.on_new_line_callback(txt, self.uid)

    def collect(self):
        with io.open(self.out_file, "r") as reader:
            log = reader.read()

        m = re.search("(?<=\n).*(?=\n$)", log)
        if m:
            user_msg = m.group(0)
        else:
            user_msg = "No match"

        if self.perf_enabled and self.is_waitable: # only for clients
            m = re.search("(?<=Performance summary:\n).*" \
                + "(?=((FATAL)|(ERROR)|(WARN)|(INFO)|(DEBUG)|(TRACE)))",
                log, re.DOTALL)
            if m:
                self.perf_file = os.path.join(self.log_dir,
                                              self.friendly_name +
                                              "_" +
                                              str(self.uid) + "_perf.log")
                with io.open(self.perf_file, "w") as writer:
                    writer.write(m.group(0))

            pattern = re.compile(r"RAWLatencyMicro\s+\d+\s+Time\s+\d+")
            self.latencies = []
            for m in re.finditer(pattern, log):
                values = m.group(0).split(' ')
                self.latencies.append((int(values[1]), int(values[3])))

        res = ProcessOutput("{0} {1}".format(self.friendly_name, self.params),
                            self.process.returncode, user_msg, self.uid,
                            self.latencies)
        return res

    def start(self):
        self.thread = threading.Thread(target=self.run)
        self.thread.start()

    def run(self):
        cmd = [self.executable]
        cmd.extend(self.params)

        g_logger.debug("Run command: {}".format(cmd))

        self.out_file = os.path.join(
            self.log_dir, self.friendly_name + "_" + str(self.uid) + "_out.log")
        os.chdir(self.log_dir)

        FNULL = open(os.devnull, 'w')
        with io.open(self.out_file, "wb", 8196) as out_writer, \
        io.open(self.out_file, "br", 8196) as out_b_reader, \
        io.open(self.out_file, "r", 8196) as out_t_reader:
            self.process = subprocess.Popen(cmd,
                                        bufsize=8196,
                                        stdout=out_writer,
                                        stderr=FNULL)
            # output only client messages
            if self.is_waitable:
                while self.process.poll() is None:
                    t = out_t_reader.readline()
                    self.invoke_on_new_line(t)
                    if not self.no_console:
                        sys.stdout.buffer.write(out_b_reader.read())
                # read leftovers
                if not self.no_console:
                    sys.stdout.buffer.write(out_b_reader.read())
                lines = out_t_reader.readlines()
                for line in lines:
                    self.invoke_on_new_line(line)

    def wait(self):
        self.thread.join()

    def stop(self):
        self.process.terminate()
        self.process.communicate()

    def kill(self):
        self.process.kill()
        self.process.communicate()


class TestConfig:
    iterations_done = 0
    leader_killed = False
    iterations_counter_lock = threading.Lock()
    config_run_timer = None
    config_last_update_time = None

    def __init__(self, bft, cli, env, simple_client_params, customize_fn):
        self.num_of_replicas = bft.num_of_replicas
        self.num_of_clients = bft.num_of_clients
        self.num_of_faulty = bft.num_of_faulties
        self.num_of_slow = bft.num_of_slow
        self.view_change_enable = bft.test_vc
        self.view_change_timeout = cli.vc_timeout
        self.num_of_iterations = cli.num_of_iterations
        self.env = env
        self.waitables = []
        self.non_waitables = []
        self.process_outputs = []
        self.config_name = bft.name.replace(",", "_")
        self.num_of_replicas_to_launch = bft.num_of_replicas_to_run
        self.on_new_line_callback = self.on_new_line_vc if bft.test_vc else \
            self.on_new_line

        self.config_timeout = 60000 + (0 if not bft.test_vc else cli.vc_timeout * 5)
        self.customize_fn = customize_fn
        self.log_dir = None
        self.simple_client_params = simple_client_params
        self.measure_perf = cli.measure_perf

    @staticmethod
    def reset_globals():
        TestConfig.iterations_done = 0
        TestConfig.leader_killed = False
        TestConfig.config_run_timer = None
        TestConfig.config_last_update_time = None

    def prepare(self, cli_args):
        exec_dir = os.path.abspath(self.env.binary_dir)
        time_stamp = time.strftime("%Y%m%dT%H%M%S")
        log_a_path = os.path.abspath(self.env.log_dir)
        self.log_dir = os.path.join(log_a_path, time_stamp)

        if not os.path.exists(self.log_dir):
            os.mkdir(self.log_dir)

        log_file_handler = logging.FileHandler(os.path.join(
                self.log_dir, "test_log.txt"))
        log_file_handler.setFormatter(g_log_formatter)
        g_logger.addHandler(log_file_handler)

        with io.open(os.path.join(
                self.log_dir, self.config_name.replace(",","_")),"w") as w:
            w.write("cli args: " + cli_args)

        # create replica's keys
        cmd = [generate_concord_keys,
               "-n", str(self.num_of_replicas),
               "-f", str(self.num_of_faulty),
               "-o", os.path.join(self.log_dir, "private_replica_")]
        subprocess.check_call(cmd)

        # create TLS certificates
        cmd = [os.path.join(log_dir, "create_tls_certs.sh"), str(self.num_of_replicas +
                                            self.num_of_clients),
               os.path.join(self.log_dir, "certs")]
        subprocess.call(cmd)

        for i in range(0, self.num_of_replicas_to_launch):
            params = ["-id", str(i),
                      "-r", str(self.num_of_replicas),
                      "-c", str(self.num_of_clients)]
            if self.view_change_enable:
                params.extend(["-vc", "-vct", str(self.view_change_timeout)])

            tp = TestProcess(self.env.replica_exec,
                             os.path.join(exec_dir, self.env.replica_exec),
                             params,
                             self.log_dir,
                             False,
                             i,
                             self.env.no_console,
                             self.measure_perf)
            self.non_waitables.append(tp)

        for i in range(0, self.num_of_clients):
            params = ["-i", str(self.num_of_iterations),
                      "-r", str(self.num_of_replicas),
                      "-cl", str(self.num_of_clients),
                      "-id", str(self.num_of_replicas + i),
                      "-f", str(self.num_of_faulty),
                      "-c", str(self.num_of_slow),
                      "-irt", str(
                    self.simple_client_params.initial_retry_timeout_milli),
                      "-minrt", str(
                    self.simple_client_params.min_retry_timeout_milli),
                      "-maxrt", str(
                    self.simple_client_params.max_retry_timeout_milli),
                      "-srft", str(
                    self.simple_client_params.sends_request_to_all_replicas_first_thresh),
                      "-srpt", str(
                    self.simple_client_params.sends_request_to_all_replicas_period_thresh),
                      "-prt", str(
                    self.simple_client_params.periodic_reset_thresh)]
            if self.measure_perf:
                params.append("-p")

            tp = TestProcess(self.env.client_exec,
                             os.path.join(exec_dir, self.env.client_exec),
                             params,
                             self.log_dir,
                             True,
                             self.num_of_replicas + i,
                             self.env.no_console,
                             self.on_new_line_callback)
            self.waitables.append(tp)

    def print_name(fn):
        def wrapper(self):
            global g_logger
            g_logger.info('Start config "{}"'.format(self.config_name))
            fn(self)
            g_logger.info('End config "{}"'.format(self.config_name))
        return wrapper

    def kill_replica(self, id):
        self.non_waitables[id].kill()
        del self.non_waitables[id]

    def kill_all(self):
        for p in self.waitables + self.non_waitables:
            p.kill()

    # this method is used to determine total number of iterations been proceed
    # so far by clients and to kill primary when needed
    def on_new_line(self, text, instanceId):
        if text is None:
            return
        pattern = re.compile(r"(?<=^Iterations count:\s)\d+")
        TestConfig.iterations_counter_lock.acquire()
        found = False
        for m in re.finditer(pattern, text):
            found = True
            TestConfig.iterations_done += int(m.group(0))
            g_logger.info("From {0}, num of iterations done: {1}".format(
                            instanceId,
                            TestConfig.iterations_done))
        if found:
            TestConfig.config_last_update_time=int(round(time.time()*1000))
        TestConfig.iterations_counter_lock.release()
        return self.on_new_line

    def on_new_line_vc(self, text, instanceId):
        retval = self.on_new_line_vc
        self.on_new_line(text, instanceId)
        TestConfig.iterations_counter_lock.acquire()
        if TestConfig.iterations_done >= int(self.num_of_iterations) * \
                int(self.num_of_clients) / 2 and not TestConfig.leader_killed:
            g_logger.info("From: {}, killing replica 0".format(instanceId))
            self.kill_replica(0)
            TestConfig.leader_killed = True
            g_logger.info("Killed replica 0")
            retval = self.on_new_line
        TestConfig.iterations_counter_lock.release()
        return retval

    def test_timer_handler(self):
        millis = int(round(time.time() * 1000))
        TestConfig.iterations_counter_lock.acquire()
        last_update_time = TestConfig.config_last_update_time
        TestConfig.iterations_counter_lock.release()
        if last_update_time is not None and \
                millis-last_update_time > self.config_timeout:
            g_logger.error("Timeout, killing all...")
            self.kill_all()
            self.collect()
            g_logger.error("Timeout, all killed")
        else:
            if last_update_time is not None and \
                    millis - last_update_time > 10000:
                g_logger.debug("No data from clients in last 10 seconds")
            TestConfig.config_timer=threading.Timer(10, self.test_timer_handler)
            TestConfig.config_timer.start()

    @print_name
    def run(self):
        curr_dir = os.getcwd()
        os.chdir(self.log_dir)
        for tp in self.non_waitables + self.waitables:
            tp.start()

        g_logger.debug("processes started")
        self.test_timer_handler()

        for tp in self.waitables:
            tp.wait()
            TestConfig.config_timer.cancel()
            res = tp.collect()
            self.process_outputs.append(res)
        for tp in self.non_waitables:
            tp.stop()
            res = tp.collect()
            self.process_outputs.append(res)

        g_logger.debug("processes stopped")
        os.chdir(curr_dir)

    def collect(self):
        if self.customize_fn:
            self.customize_fn(self)
        return self.process_outputs

    @staticmethod
    def create(env, cli, bft, simple_client_params, customize_fn):
        return TestConfig(bft, cli, env, simple_client_params, customize_fn)


def parse_client_histogram(file):
    hist = None
    with io.open(file, "r") as reader:
        hist = reader.read()

    ptrn = re.compile(r"((?<=[\[\s*])\d*(?=,)),\s*((?<=[,\s*])\d*(?=\s*\)))\s*\)\s*((?<=[" \
                      "\s*])\d*(?=\s*))")
    res = []
    for (low, high, count) in re.findall(ptrn, hist):
        res.append((int(low), int(high), int(count)))
    return res


def process_output(process_outputs, perf_enabled):
    res = ""
    exit_code = 0
    if perf_enabled:
        perf_hist = histogram.Histogram()
        clients_perf_info = []

    count = 0
    for output in process_outputs:
        res += "{0} is done with exit code {1} and message '{2}'".format(
            output.exec_name,
            output.exit_code,
            output.user_msg)
        res += os.linesep
        exit_code += output.exit_code

        begin_milli_time = 0xFFFFFFFFFFFFFFFF
        end_milli_time = 0
        num_of_ops = 0
        if perf_enabled and output.latencies:
            for l in output.latencies:
                tmp = l[1] - (l[0] / 1000)
                if begin_milli_time > tmp:
                    begin_milli_time = tmp
                if end_milli_time < l[1]:
                    end_milli_time = l[1]
                num_of_ops += 1
                perf_hist.add(l[0])
            duration_milli = end_milli_time - begin_milli_time
            clients_perf_info.append((duration_milli, num_of_ops))
            count += 1

    w_tp = 0
    sum_dur = 0
    sum_ops = 0
    for i in range(0, count):
        w_tp += clients_perf_info[i][1] / clients_perf_info[i][0] * clients_perf_info[i][1]
        sum_ops += clients_perf_info[i][1]
        sum_dur += clients_perf_info[i][0]
    if sum_ops > 0:
        w_tp /= sum_ops

    g_logger.debug("output: {0}, {1}".format(res, exit_code))
    if perf_enabled:
        buf = io.StringIO()
        buf.write("Performance summary:\n")
        buf.write("Average duration(sec): %.2f" % (sum_dur / count / 1000))
        buf.write(", Average throughput(ops/sec): %.2f\n" % (w_tp * 1000))
        buf.write(perf_hist.to_string())
        t = buf.getvalue()
    else:
        t = None

    return TestResult(exit_code, res, t)


def customize_test_vc_never_ends(test_config):
    for o in test_config.process_outputs:
        o.exit_code = 0
        o.user_message = "Success"


def main():
    parser = argparse.ArgumentParser(
        description="Simple test automation script",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-bd",
                        help="Folder with BFT executables",
                        default=bin_dir)
    parser.add_argument("-ld",
                        help="Directory for log files",
                        default=log_dir)
    parser.add_argument("-nc",
                        help="Do not print processes' output to the console",
                        action="store_true")
    parser.add_argument("-l",
                        help="Set log level for this script (not for the " +
                        "executables)",
                        default="INFO")
    parser.add_argument("-i",
                        help="Set number of operations to run",
                        default=2800,type=int)
    parser.add_argument("-vct",
                        default=60000,
                        type=int,
                        help="View change timeout (if VC is enabled)")
    parser.add_argument("-p",
                        action="store_true",
                        help="Enable perfomance measurement by test client")
    parser.add_argument("-bft",
                       default=["n=3,r=3,f=1,c=0,cl=1"],
                       nargs="+",
                       help="List of BFT configuration sets to use. All " +
                       "parameters should be separated by ','. " +
                       "n - total number of " + "replicas, r - number of " +
                       "running replicas, " +
                       "f - max. number of byzantine replicas, " +
                       "c - max. number of slow replicas, " +
                       "cl - number of running clients, " +
                       "testViewChange - if present, viewChange will be " +
                       "triggered after 1/2 of iterations done. " +
                       "Please note, if View change is enabled AND n - r " +
                       "<= f - the system will not progress after primary " +
                       "is killed, and the test will SUCCESS after timeout " +
                       "- this is desired behaviour for this configuration." +
                       "You can provide set " +
                       "of configurations in the form of " +
                       "set1 set2 .... setN - " +
                       "each of them will be ran as separate test")

    args = parser.parse_args()
    env = Environment("server", "client", args.bd, args.ld, args.nc)

    global g_logger
    g_logger.setLevel(args.l)

    configs = []
    for param in args.bft:
        set = param.split(",")
        n=r=f=c=cl=None
        vc=False
        for p_expr in set:
            p = p_expr.split("=")
            if p[0] == "n":
                n = int(p[1])
            elif p[0] == "r":
                r = int(p[1])
            elif p[0] == "f":
                f = int(p[1])
            elif p[0] == "c":
                c = int(p[1])
            elif p[0] == "cl":
                cl = int(p[1])
            elif p[0] == "testViewChange":
                vc = True
            else:
                print("Unsupported parameter: {}".format(p[0]))
                return -1
        if n is None or r is None or f is None or c is None or cl is None:
            print("Unable to parse BFT parameters")
            return -1
        if 2 * f + 1 != n:
            g_logger.error("N = 2f + 1 is not satisfied")
            return -1
        max_r = max(3, multiprocessing.cpu_count() * 3)
        if r > max_r or r > 63:
            g_logger.error("max number of running replicas on your machine is "
                           "" + str(max_r))
            return -1
        if r < f + 1:
            g_logger.error("r >= f + 1 is not satisfied")
            return -1
        if r > n:
            g_logger.error("r <= n is not satisfied")
            return -1
        if cl > 10:
            g_logger.error("Max number of clients is 10")
            return -1

        configs.append(BFT(param, n, r, f, c, cl, vc))

    if args.i < 200 or args.i > 10000:
        g_logger.error("-i value should be between 200 and 10000")
        return -1

    cli = CLI(args.i, args.vct, args.p)
    scp = SimpleClientParams()

    for cf in configs:
        g_logger.debug("Creating test configuration {0} with n={1},r={2} \
f = {3}, c = {4}, testViewChange = {5}, cl = {6}".format(
                           cf.name, cf.num_of_replicas,
                           cf.num_of_replicas_to_run,
                           cf.num_of_faulties, cf.num_of_slow, cf.test_vc, cf.num_of_clients))

        c = TestConfig.create(env, cli, cf, scp,
              customize_test_vc_never_ends if cf.vc_will_fail else None)

        TestConfig.reset_globals()
        c.prepare(str(sys.argv))
        c.run()
        out = c.collect()
        result = process_output(out, cli.measure_perf)
        if result.exit_code != 0:
            g_logger.info("CONFIGURATION RUN FAIL")
            # THIS IS DONE ON PURPOSE TO SUPPORT AUTOMATIC TESTING
            print("TESTS FAIL")
            return result.exit_code
        else:
            # THIS IS DONE ON PURPOSE TO SUPPORT AUTOMATIC TESTING
            print("CONFIGURATION RUN SUCCESS")
            g_logger.info("CONFIGURATION RUN SUCCESS")

        if result.perf_summary:
            g_logger.info(result.perf_summary)

    g_logger.info("TEST SUCCESS")
    print("DONE")
    return 0

main()
