#!/usr/bin/env python
#-*- coding: utf-8 -*-
from multiprocessing import cpu_count
from subprocess import Popen, check_call
import os
import shutil
import argparse
import logging
import time

def run_perf_test(cmd, xmls_path, output_folder):
    output_path = os.path.join(output_folder, "perf_stress_run.txt")
    f = open(output_path, 'w')
    p = Popen("{} --skip-tags=long --recursive --input-files {}".format(cmd, xmls_path), shell=True, stdout=f, stderr=f)
    return p

def run_func_test(cmd, output_prefix, num_processes, skip_tests_option):
    output_paths = [os.path.join(output_prefix, "stress_test_run_{}.txt".format(i)) for i in range(num_processes)]
    f = open(output_paths[0], 'w')
    main_command = "{} {}".format(cmd, skip_tests_option)
    logging.info("Run func tests main cmd '%s'", main_command)
    pipes = [Popen(main_command, shell=True, stdout=f, stderr=f)]
    for output_path in output_paths[1:]:
        time.sleep(0.5)
        f = open(output_path, 'w')
        full_command = "{} --order=random {}".format(cmd, skip_tests_option)
        logging.info("Run func tests '%s'", full_command)
        p = Popen(full_command, shell=True, stdout=f, stderr=f)
        pipes.append(p)
    return pipes

def check_clickhouse_alive(cmd):
    try:
        logging.info("Checking ClickHouse still alive")
        check_call("{} --query \"select 'Still alive'\"".format(cmd), shell=True)
        return True
    except:
        return False

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
    parser = argparse.ArgumentParser(description="ClickHouse script for running stresstest")
    parser.add_argument("--test-cmd", default='clickhouse-test')
    parser.add_argument("--skip-func-tests", default='')
    parser.add_argument("--client-cmd", default='clickhouse-client')
    parser.add_argument("--perf-test-cmd", default='clickhouse-performance-test')
    parser.add_argument("--perf-test-xml-path", default='/usr/share/clickhouse-test/performance/')
    parser.add_argument("--server-log-folder", default='/var/log/clickhouse-server')
    parser.add_argument("--output-folder")
    parser.add_argument("--num-parallel", default=cpu_count() // 3);

    args = parser.parse_args()
    func_pipes = []
    perf_process = None
    try:
        perf_process = run_perf_test(args.perf_test_cmd, args.perf_test_xml_path, args.output_folder)
        func_pipes = run_func_test(args.test_cmd, args.output_folder, args.num_parallel, args.skip_func_tests)

        logging.info("Will wait functests to finish")
        while True:
            retcodes = []
            for p in func_pipes:
                if p.poll() is not None:
                    retcodes.append(p.returncode)
            if len(retcodes) == len(func_pipes):
                break
            logging.info("Finished %s from %s processes", len(retcodes), len(func_pipes))
            time.sleep(5)

        if not check_clickhouse_alive(args.client_cmd):
            raise Exception("Stress failed, results in logs")
        else:
            logging.info("Stress is ok")
    except Exception as ex:
        raise ex
    finally:
        if os.path.exists(args.server_log_folder):
            logging.info("Copying server log files")
            for log_file in os.listdir(args.server_log_folder):
                shutil.copy(os.path.join(args.server_log_folder, log_file), os.path.join(args.output_folder, log_file))