#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import copy
import multiprocessing
import os
import subprocess
import time
from argparse import Namespace as Args
from subprocess import Popen
from typing import Dict, List, Optional

from test_config import LARGE_DATASET, SMALL_DATASET, DatabaseMode, DatasetConstants

# paths
SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
BASE_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", ".."))
BUILD_DIR = os.path.join(BASE_DIR, "build")
MEASUREMENTS_FILE = os.path.join(SCRIPT_DIR, ".apollo_measurements")
KEY_FILE = os.path.join(SCRIPT_DIR, ".key.pem")
CERT_FILE = os.path.join(SCRIPT_DIR, ".cert.pem")

# long running stats file
STATS_FILE = os.path.join(SCRIPT_DIR, ".long_running_stats")

# get number of threads
THREADS = os.environ["THREADS"] if "THREADS" in os.environ else multiprocessing.cpu_count()


def generate_temporary_ssl_certs():
    # https://unix.stackexchange.com/questions/104171/create-ssl-certificate-non-interactively
    subj = "/C=HR/ST=Zagreb/L=Zagreb/O=Memgraph/CN=db.memgraph.com"
    subprocess.run(
        [
            "openssl",
            "req",
            "-new",
            "-newkey",
            "rsa:4096",
            "-days",
            "365",
            "-nodes",
            "-x509",
            "-subj",
            subj,
            "-keyout",
            KEY_FILE,
            "-out",
            CERT_FILE,
        ],
        check=True,
    )


def remove_certificates() -> None:
    os.remove(KEY_FILE)
    os.remove(CERT_FILE)


def parse_arguments() -> Args:
    # parse arguments
    parser = argparse.ArgumentParser(description="Run stress tests on Memgraph.")
    parser.add_argument("--memgraph", default=os.path.join(BUILD_DIR, "memgraph"))
    parser.add_argument("--log-file", default="")
    parser.add_argument("--data-directory", default="")
    parser.add_argument("--python", default=os.path.join(SCRIPT_DIR, "ve3", "bin", "python3"), type=str)
    parser.add_argument("--large-dataset", action="store_const", const=True, default=False)
    parser.add_argument("--small-dataset", action="store_const", const=True, default=False)
    parser.add_argument("--specific-test", default="")
    parser.add_argument("--use-ssl", action="store_const", const=True, default=False)
    parser.add_argument("--verbose", action="store_const", const=True, default=False)
    return parser.parse_args()


def wait_for_server(port, delay=0.1) -> None:
    cmd = ["nc", "-z", "-w", "1", "127.0.0.1", str(port)]
    while subprocess.call(cmd) != 0:
        time.sleep(0.01)
    time.sleep(delay)


def start_memgraph(args: Args, memgraph_options: List[str]) -> Popen:
    """Starts Memgraph and return the process"""
    cwd = os.path.dirname(args.memgraph)
    cmd = [
        args.memgraph,
        "--bolt-num-workers=" + str(THREADS),
        "--storage-properties-on-edges=true",
        "--storage-snapshot-on-exit=true",
        "--storage-snapshot-interval-sec=600",
        "--storage-snapshot-retention-count=1",
        "--storage-wal-enabled=true",
        "--storage-recover-on-startup=false",
        "--query-execution-timeout-sec=1200",
        "--bolt-server-name-for-init=Neo4j/",
    ] + memgraph_options
    if not args.verbose:
        cmd += ["--log-level", "WARNING"]
    if args.log_file:
        cmd += ["--log-file", args.log_file]
    if args.data_directory:
        cmd += ["--data-directory", args.data_directory]
    if args.use_ssl:
        cmd += ["--bolt-cert-file", CERT_FILE, "--bolt-key-file", KEY_FILE]
    memgraph_proc = subprocess.Popen(cmd, cwd=cwd)
    wait_for_server(7687)

    assert memgraph_proc.poll() is None, "The database binary died prematurely!"

    return memgraph_proc


def stop_memgraph(proc_mg: Popen) -> None:
    proc_mg.terminate()
    ret_mg = proc_mg.wait()
    if ret_mg != 0:
        raise Exception("Memgraph binary returned non-zero ({})!".format(ret_mg))


def run_test(args: Args, test: str, options: List[str], timeout: int, mode: DatabaseMode) -> float:
    """Runs tests for a set of specific database configuration.

    Args:
        args: Arguments passed to the test
        test: Test name
        options: List of options specific for each test
        timeout: Timeout in minutes
        mode: DatabaseMode (storage mode & isolation level pair)
    """
    print("Running test '{}'".format(test))

    binary = _find_test_binary(args, test)

    # start test
    cmd = (
        binary
        + [
            "--worker-count",
            str(THREADS),
            "--isolation-level",
            mode.isolation_level,
            "--storage-mode",
            mode.storage_mode,
        ]
        + options
    )
    start = time.time()
    ret_test = subprocess.run(cmd, cwd=SCRIPT_DIR, timeout=timeout * 60)

    if ret_test.returncode != 0:
        raise Exception("Test '{}' binary returned non-zero ({})!".format(test, ret_test.returncode))

    runtime = time.time() - start
    print("    Done after {:.3f} seconds".format(runtime))

    return runtime


def _find_test_binary(args: Args, test: str) -> List[str]:
    if test.endswith(".py"):
        logging = "DEBUG" if args.verbose else "WARNING"
        return [args.python, "-u", os.path.join(SCRIPT_DIR, test), "--logging", logging]

    if test.endswith(".cpp"):
        exe = os.path.join(BUILD_DIR, "tests", "stress", test[:-4])
        return [exe]

    raise Exception("Test '{}' binary not supported!".format(test))


def run_stress_test_suite(args: Args) -> Optional[Dict[str, float]]:
    def cleanup(memgraph_proc):
        if memgraph_proc.poll() != None:
            return
        memgraph_proc.kill()
        memgraph_proc.wait()

    runtimes = {}

    if args.large_dataset and args.small_dataset:
        raise Exception("Choose only a large dataset or a small dataset for stress testing!")

    dataset = LARGE_DATASET if args.large_dataset else SMALL_DATASET
    if args.specific_test:
        dataset = [x for x in dataset if x[DatasetConstants.TEST] == args.specific_test]
        if not len(dataset):
            raise Exception("Specific dataset is not found for stress testing!")

    for test in dataset:
        if args.use_ssl:
            test[DatasetConstants.OPTIONS] += ["--use-ssl"]

        for mode in test[DatasetConstants.MODE]:
            test_run = copy.deepcopy(test)
            # Run for every specified combination of storage mode, serialization mode, etc (if extended)
            test_run[DatasetConstants.MODE] = mode
            try:
                memgraph_options = (
                    test_run[DatasetConstants.MEMGRAPH_OPTIONS] if DatasetConstants.MEMGRAPH_OPTIONS in test_run else []
                )
                if DatasetConstants.MEMGRAPH_OPTIONS in test_run:
                    del test_run[DatasetConstants.MEMGRAPH_OPTIONS]
                memgraph_proc = start_memgraph(args, memgraph_options)
            except Exception as ex:
                print("Exception occured while starting memgraph", ex)
                return None
            err = False
            try:
                runtime = run_test(args, **test_run)
                runtimes[os.path.splitext(test[DatasetConstants.TEST])[0]] = runtime
            except Exception as ex:
                print(f"Failed to execute {test_run} with following exception:", ex)
                err = True
            finally:
                stop_memgraph(memgraph_proc)
                cleanup(memgraph_proc)
            if err:
                return None

    return runtimes


def write_stats(runtimes: Dict[str, float]) -> None:
    measurements = ""
    for key, value in runtimes.items():
        measurements += "{}.runtime {}\n".format(key, value)
    with open(STATS_FILE) as f:
        stats = f.read().split("\n")
    measurements += "long_running.queries.executed {}\n".format(stats[0])
    measurements += "long_running.queries.failed {}\n".format(stats[1])
    with open(MEASUREMENTS_FILE, "w") as f:
        f.write(measurements)


if __name__ == "__main__":
    args = parse_arguments()

    if args.use_ssl:
        generate_temporary_ssl_certs()

    runtimes = run_stress_test_suite(args)
    if runtimes is None:
        print("Some stress tests have failed")
        exit(1)
    assert runtimes is not None
    if args.use_ssl:
        remove_certificates()
    write_stats(runtimes)

    print("Successfully ran stress tests!")
