#!/usr/bin/python3 -u
import argparse
import json
import os
import subprocess
import sys
import tempfile
import time

SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
PROJECT_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", "..", ".."))


def execute_test(**kwargs):
    client_binary = kwargs.pop("client")
    server_binary = kwargs.pop("server")
    storage_directory = kwargs.pop("storage")

    start_server = kwargs.pop("start_server", True)
    endpoint = kwargs.pop("endpoint", "")
    interval = kwargs.pop("interval", 1)
    duration = kwargs.pop("duration", 5)

    timeout = duration * 2 if "hang" not in kwargs else duration * 2 + 60
    success = False

    server_args = [server_binary, "--interval", interval,
                   "--duration", duration]
    for flag, value in kwargs.items():
        flag = "--" + flag.replace("_", "-")
        # We handle boolean flags here. The type of value must be `bool`, and
        # the value must be `True` to supply a boolean flag (a flag that only
        # has --flag, without the value).
        if value is True:
            server_args.append(flag)
        else:
            server_args.extend([flag, value])

    client_args = [client_binary, "--interval", interval,
                   "--duration", duration,
                   "--storage-directory", storage_directory]
    if endpoint:
        client_args.extend(["--endpoint", endpoint])

    server = None
    if start_server:
        server = subprocess.Popen(list(map(str, server_args)))
        time.sleep(0.1)
        assert server.poll() is None, "Server process died prematurely!"

    try:
        subprocess.run(list(map(str, client_args)), timeout=timeout,
                       check=True)
    finally:
        if server is None:
            success = True
        else:
            server.terminate()
            try:
                success = server.wait(timeout=5) == 0
            except subprocess.TimeoutExpired:
                server.kill()
    return success


TESTS = [
    {},
    {"interval": 2},
    {"duration": 10},
    {"interval": 2, "duration": 10},
    {"redirect": True},
    {"no_response_count": 2},
    {"wrong_code_count": 2},
    {"hang": True, "duration": 0},
    {"path": "/nonexistant/", "no_check": True},
    {"endpoint": "http://127.0.0.1:9000/nonexistant/", "no_check": True},
    {"start_server": False},
    {"startups": 4, "no_check_duration": True},  # the last 3 tests failed
                                                 # to send any data + this test
    {"add_garbage": True}
]

if __name__ == "__main__":
    server_binary = os.path.join(SCRIPT_DIR, "server.py")
    client_binary = os.path.join(PROJECT_DIR, "build", "tests",
                                 "integration", "telemetry", "client")
    if not os.path.exists(client_binary):
        client_binary = os.path.join(PROJECT_DIR, "build_debug", "tests",
                                     "integration", "telemetry", "client")
    kvstore_console_binary = os.path.join(PROJECT_DIR, "build", "tests",
                                          "manual", "kvstore_console")
    if not os.path.exists(kvstore_console_binary):
        kvstore_console_binary = os.path.join(PROJECT_DIR, "build_debug",
                                              "tests", "manual",
                                              "kvstore_console")

    parser = argparse.ArgumentParser()
    parser.add_argument("--client", default=client_binary)
    parser.add_argument("--server", default=server_binary)
    parser.add_argument("--kvstore-console", default=kvstore_console_binary)
    args = parser.parse_args()

    storage = tempfile.TemporaryDirectory()

    for test in TESTS:
        print("\033[1;36m~~ Executing test with arguments:",
              json.dumps(test, sort_keys=True), "~~\033[0m")

        if test.pop("add_garbage", False):
            proc = subprocess.Popen([args.kvstore_console, "--path",
                                     storage.name], stdin=subprocess.PIPE,
                                    stdout=subprocess.DEVNULL)
            proc.communicate("put garbage garbage".encode("utf-8"))
            assert proc.wait() == 0

        try:
            success = execute_test(client=args.client, server=args.server,
                                   storage=storage.name, **test)
        except Exception as e:
            print("\033[1;33m", e, "\033[0m", sep="")
            success = False

        if not success:
            print("\033[1;31m~~", "Test failed!", "~~\033[0m")
            sys.exit(1)
        else:
            print("\033[1;32m~~", "Test ok!", "~~\033[0m")

    sys.exit(0)