#!/usr/bin/env python3
# -*- coding:utf-8 -*-
"""
Copyright (c): 2021, Huawei Tech. Co., Ltd.
Description  : A utility to start/stop/restart gds or roach_client.
"""
try:
    import os
    import sys
    import getopt
    import getpass
    import subprocess

    sys.path.append(os.path.split(os.path.realpath(__file__))[0] + "/../../script/")
    from gspylib.common.ErrorCode import ErrorCode
    from gspylib.common.GaussLog import GaussLog
    from common_tools import GDSUtils

except ImportError as e:
    sys.exit("[GAUSS-52200] : Unable to import module: %s." % str(e))

# APP_NAME: gds or roach_client
APP_NAME = __file__[(0 if __file__.rfind("/") == -1 else __file__.rfind("/") + 1):__file__.rfind("_")]

# global variables
install_dir = ""
gds_user = ""
install_hosts = []
oper_dict = {}
curr_oper = ""
curr_conf = None
local_flag = False
remote_hosts = []

# roach_client only can recongnize less.
if APP_NAME == "gds":
    param_dict = {"-d": "data_dir", "-e": "err_dir", "-E": "err_seg",
                  "-R": "log_seg", "-S": "data_seg", "-H": "secure_ip_range",
                  "-t": "parallel", "-l": "log_file", "-r": "recursive",
                  "-D": "daemon", "-s": "status_file",
                  "--pipe-timeout": "pipe_timeout",
                  "--debug-level": "debug_level",
                  "--enable-ssl": "enable_ssl",
                  "--ssl-dir": "ssl_dir",

                  }
else:
    param_dict = {"-R": "log_seg", "-H": "secure_ip_range", "-l": "log_file",
                  "-D": "daemon",
                  "--debug-level": "debug_level",
                  "--enable-ssl": "enable_ssl",
                  "--ssl-dir": "ssl_dir",
                  }


class StartupConf:
    def __init__(self):
        self.port = -1
        self.name = None
        self.log_file = None
        self.data_dir = None
        self.err_dir = None
        self.data_seg = None
        self.err_seg = None
        self.log_seg = None
        self.ctl_file = None
        self.recursive = False
        self.daemon = False
        self.enable_ssl = False
        self.ssl_dir = None
        self.secure_ip_range = None
        self.ip = "0.0.0.0"
        self.pipe_timeout = None
        self.debug_level = None
        self.parallel = -1
        self.help = False
        self.status_file = None

        self.error_msg = ""

    def is_valid(self):
        if self.help:
            return True
        if self.port < 1024 or self.port > 65535:
            self.error_msg = ErrorCode.GAUSS_500['GAUSS_50004'] % 'p' + \
                             "Port number should range between 1024 and 65536."
            return False
        if self.data_dir is None:
            self.error_msg = ErrorCode.GAUSS_500['GAUSS_50014'] % '-d'
            return False
        if self.secure_ip_range is None:
            self.error_msg = ErrorCode.GAUSS_500['GAUSS_50014'] % '-H'
            return False
        return True


def usage():
    """
%s_ctl is a utility to start, stop or restart %s servers.
Usage:
    %s_ctl start --host [/path/to/hostfile | ipaddr1,ipaddr2...] -p PORT -d DATADIR -H ALLOW_IPs [%s other options]
    %s_ctl stop --host [/path/to/hostfile | ipaddr1,ipaddr2...] -p PORT
    %s_ctl restart --host [/path/to/hostfile | ipaddr1,ipaddr2...] -p PORT
Common options:
    -p                      Almost same as %s's '-p' option, except here just a port number.
    --host                  The path of ip address list for nodes, or ip address string seprated by comma.
    --help                  Show help information.
    -V                      Show version information.
Original %s options:
%s
    """
    if APP_NAME == 'gds':
        app_special_options_docstr = """\
    -d dir                  Set data directory.
    -l log_file             Set log file.
    -H secure_ip_range      Set secure IP checklist in CIDR notation. Required for GDS to start.
    -e dir                  Set error log directory.
    -E size                 Set size of per error log segment.(0 < size < 1TB)
    -S size                 Set size of data segment.(1MB < size < 100TB)
    -R size                 Set size of log rotate.(1MB < size < 1TB) default 16M
    -t worker_num           Set number of worker thread in multi-thread mode, the upper limit is 200.\
     If without setting, the default value is 8.
    --enable-ssl            Set GDS connection with SSL security.
    --ssl-dir               Set the directory holding certificate files.
    --debug-level           Set the debug level, which decides how much debug information we print in log.\
     Default value is 0.
                                    0:  OFF. Keep the minimum amount of log. Only contains session descriptions.
                                    1:  NORMAL. Log the basic "connnect" and "disconnect" information.\
                                     Useful for identify issues with network traffics.
                                    2:  ON. Keep detailed log about node connections, package sent and received,\
                                     status information.
    --pipe-timeout          Set the pipe timeout for no landing export/import (>1s). Default 1h/60m/3600s.
    -s status_file          Enable GDS status report.
    -D                      Run the GDS as a daemon process.
    -r                      Read the working directory recursively.
"""
    elif APP_NAME == "roach_client":
        app_special_options_docstr = """\
    -l log_file             Set log file.
    -H secure_ip_range      Set secure IP checklist in CIDR notation. Required for ROACH_CLIENT to start.
    -R size                 Set size of log rotate.(1MB < size < 1TB) default 16M
    --enable-ssl            Set ROACH_CLIENT connection with SSL security.
    --ssl-dir               Set the directory holding certificate files.
    --debug-level           Set the debug level, which decides how much debug information we print in log.\
     Default value is 0.
                                    0:  OFF. Keep the minimum amount of log. Only contains session descriptions.
                                    1:  NORMAL. Log the basic "connnect" and "disconnect" information.\
                                     Useful for identify issues with network traffics.
                                    2:  ON. Keep detailed log about node connections, package sent and received,\
                                     status information.
    -D                      Run the ROACH_CLIENT as a daemon process.
"""
    else:
        app_special_options_docstr = "    None."
    print(usage.__doc__ % ((APP_NAME,) * 8 + (app_special_options_docstr,)))


def parse_command_line():
    """
    function: Check parameter from command line
    """
    global curr_oper
    global curr_conf
    global install_hosts

    try:
        if (len(sys.argv) == 1):
            GaussLog.printMessage("[GAUSS-50001] : Incorrect parameter."
                                  " Operation type is required: start, stop or restart.")
            usage()
            sys.exit(1)
        elif (sys.argv[1] == "--help"):
            usage()
            sys.exit(0)
        elif (sys.argv[1] == "-V"):
            GDSUtils.get_version_info(APP_NAME)
            sys.exit(0)
        elif (sys.argv[1] not in oper_dict.keys()):
            GaussLog.exitWithError(ErrorCode.GAUSS_500['GAUSS_50011'] % ("operation type", sys.argv[1]))
        curr_oper = sys.argv[1]
        opts, args = getopt.getopt(sys.argv[2:], "d:e:E:p:R:S:t:H:l:rDV",
                                   ["host=", "pipe-timeout=", "debug-level=", "help"])

        if len(args) > 0:
            GaussLog.exitWithError(ErrorCode.GAUSS_500["GAUSS_50000"] % str(args[0]))

        curr_conf = StartupConf()
        for (key, value) in opts:
            if key == "--help":
                usage()
                sys.exit(0)
            elif key in param_dict.keys():
                curr_conf.__dict__[param_dict[key]] = value
            elif key == '-p':
                curr_conf.port = int(value)
            elif key == '-V':
                GDSUtils.get_version_info(APP_NAME)
                sys.exit()
            elif key == "--host":
                install_hosts = GDSUtils.get_hosts_from_param(value)
            else:
                GaussLog.exitWithError(ErrorCode.GAUSS_500["GAUSS_50000"] % key)
        if curr_conf.port == -1:
            curr_conf.port = 8098
        if curr_conf.daemon == "":
            curr_conf.daemon = True
        if curr_conf.recursive == "":
            curr_conf.recursive = True
        if curr_conf.enable_ssl == "":
            curr_conf.enable_ssl = True

    except Exception as e:
        GaussLog.printMessage(ErrorCode.GAUSS_500["GAUSS_50015"] % str(e))
        usage()
        sys.exit()


def check_parameters():
    """
    function: Check command line parameters, and init global variables related to those parameters
    input  : NA
    output : NA
    """
    global local_flag
    global local_ip
    global remote_hosts

    if not install_hosts:
        GaussLog.exitWithError(ErrorCode.GAUSS_500["GAUSS_50001"] % '-host')

    if curr_conf is None:
        GaussLog.exitWithError(ErrorCode.GAUSS_500["GAUSS_50009"])
    elif APP_NAME == "gds" and curr_oper == "start" and not curr_conf.is_valid():
        GaussLog.exitWithError(curr_conf.error_msg)

    if curr_conf.data_dir:
        curr_conf.data_dir = os.path.realpath(curr_conf.data_dir)
    local_ip, local_flag, remote_hosts = GDSUtils.divide_local_remote(install_hosts)
    curr_conf.ip = local_ip


def read_install_conf():
    """
    function: try to read GDS_INSTALL_DIR from OS env
    """
    global install_dir
    global logger
    if GDSUtils.ENV_NAME_GDS_INSTALL_DIR not in os.environ.keys():
        GaussLog.exitWithError(ErrorCode.GAUSS_535["GAUSS_53507"])
    install_dir = os.environ[GDSUtils.ENV_NAME_GDS_INSTALL_DIR]
    if not os.path.exists(install_dir):
        GaussLog.exitWithError(ErrorCode.GAUSS_535["GAUSS_53505"])
    logger = GaussLog("%s/logs/%s_ctl.log" % (install_dir, APP_NAME))


def is_process_exist(port):
    """
    function: check whether a process with this port has been already started.
    """
    status, output = subprocess.getstatusoutput("lsof -i :%d | grep -w %s | wc -l" % (port, APP_NAME))
    if status != 0:
        GaussLog.exitWithError(ErrorCode.GAUSS_535["GAUSS_53506"] % output)
    try:
        result = int(output)
        return result > 0
    except Exception as e:
        GaussLog.exitWithError(ErrorCode.GAUSS_535["GAUSS_53506"] % str(e))


def add_watch_dog(conf, start_cmd):
    """
    function: add an item to.
    """
    watch_cmd = "source %s; nohup $GDS_INSTALL_DIR/script/watch.sh %d %s \"%s\" &" % (GDSUtils.EXEC_ENV_FILE,
                                                                                      conf.port,
                                                                                      APP_NAME,
                                                                                      start_cmd)
    cmd = "(crontab -l; echo '* * * * * %s') | crontab" % watch_cmd
    (status, output) = subprocess.getstatusoutput(cmd)
    if status != 0:
        GaussLog.exitWithError(ErrorCode.GAUSS_535["GAUSS_53508"] % APP_NAME)


def del_watch_dog(conf):
    """
    function: delete an item from crontab.
    """
    cmd = "(crontab -l | grep -v %d) | crontab" % (conf.port)
    (status, output) = subprocess.getstatusoutput(cmd)
    if status != 0:
        GaussLog.exitWithError(ErrorCode.GAUSS_535["GAUSS_53509"] % APP_NAME)


def save_cmd_to_file(conf, start_cmd):
    """
    function: save start cmd to temp file
    """
    cmd = "echo '%s' > %s/.startcmd.%d" % (start_cmd, install_dir, conf.port)
    (status, output) = subprocess.getstatusoutput(cmd)
    if status != 0:
        GaussLog.exitWithError(ErrorCode.GAUSS_535["GAUSS_53510"] % APP_NAME)


def get_cmd_from_file(conf):
    """
    function: get start cmd to temp file
    """
    cmd = "cat %s/.startcmd.%d" % (install_dir, conf.port)
    (status, output) = subprocess.getstatusoutput(cmd)
    if status != 0:
        GaussLog.exitWithError(ErrorCode.GAUSS_535["GAUSS_53511"] % APP_NAME)
    return output


def start_gds(conf):
    """
    function: operation 'start' handler, start an app locally, dispatch to remote hosts with local mode.
              Done by gds_user.
    """

    if local_flag:
        # start app locally
        start_gds_locally(conf)
    if remote_hosts:
        try:
            new_argv = sys.argv[:]
            idx = sys.argv.index("--host")

            # an inner function

            def ssh_cmd_handler(_ip):
                new_argv[idx + 1] = _ip
                result = {}
                cmd0 = ' '.join(new_argv)
                cmd = "source %s; cd %s/script; %s" % (GDSUtils.EXEC_ENV_FILE, install_dir, cmd0)
                cmd_with_ip = "export LD_LIBRARY_PATH=/lib64:$LD_LIBRARY_PATH; ssh %s \'%s\'" % (_ip, cmd)
                (status, output) = subprocess.getstatusoutput(cmd_with_ip)
                result[_ip] = output
                return result

            parall_result = GDSUtils.exec_command_remote_parallelly(ssh_cmd_handler,
                                                                    remote_hosts,
                                                                    timeout=GDSUtils.TIMEOUT)
            for key in parall_result:
                logger.log("[host %s]\n%s" % (key, parall_result[key]))
        except Exception as e:
            logger.logExit(str(e))


def start_gds_locally(conf):
    """
    function: execute this fuction to start an app process.
    """
    if is_process_exist(conf.port):
        logger.logExit(('{0} at port [{1}] already exists.').format(APP_NAME, conf.port))

    if (conf.ctl_file is not None) and (os.path.exists(conf.ctl_file)):
        subprocess.getstatusoutput('rm {0}'.format(conf.ctl_file))
    cmd = (APP_NAME + ' -D')
    if conf.data_dir is not None:
        cmd += (' -d {0}'.format(conf.data_dir))
    if conf.port is not None:
        cmd += (' -p {0}:{1:d}'.format(conf.ip, conf.port))
    if conf.log_file is not None:
        cmd += (' -l {0}'.format(conf.log_file))
    if conf.recursive:
        cmd += ' -r'
    if conf.secure_ip_range is not None:
        cmd += (' -H {0}'.format(conf.secure_ip_range))
    if conf.parallel != -1:
        cmd += (' -t {0}'.format(conf.parallel))
    if conf.err_dir is not None:
        cmd += (' -e {0}'.format(conf.err_dir))
    if conf.log_seg is not None:
        cmd += (' -R {0}'.format(conf.log_seg))
    if conf.data_seg is not None:
        cmd += (' -S {0}'.format(conf.data_seg))
    if conf.err_seg is not None:
        cmd += (' -E {0}'.format(conf.err_seg))
    if conf.status_file is not None:
        cmd += (' -s {0}'.format(conf.status_file))
    if conf.debug_level is not None:
        cmd += (' --debug-level {0}'.format(conf.debug_level))
    if conf.pipe_timeout is not None:
        cmd += (' --pipe-timeout {0}'.format(conf.pipe_timeout))
    if conf.enable_ssl:
        cmd += ' --enable-ssl'
    if conf.ssl_dir is not None:
        cmd += (' --ssl-dir {0}'.format(conf.ssl_dir))
    (status, output) = subprocess.getstatusoutput(cmd)
    local_mark = "[Local]\n" if remote_hosts else ""
    if status != 0:
        logger.logExit("%sStart %s at port %d [ERROR]: %s" % (local_mark, APP_NAME.upper(), conf.port, output))
    add_watch_dog(conf, cmd)
    new_argv = sys.argv[:]
    idx = sys.argv.index("--host")
    new_argv[idx + 1] = conf.ip
    save_cmd_to_file(conf, " ".join(new_argv))
    logger.log("%sStart %s at port %d [OK]" % (local_mark, APP_NAME.upper(), conf.port))


def stop_gds(conf):
    """
    function: operation 'stop' handler, stop an app locally, dispatch to remote hosts with local mode.
              Done by gds_user.
              Also used for restart.
    """
    if remote_hosts:
        # dispatch to remote hosts with local_mode
        try:
            new_argv = sys.argv[:]
            idx = sys.argv.index("--host")

            # an inner function

            def ssh_cmd_handler(_ip):
                new_argv[idx + 1] = _ip
                cmd0 = ' '.join(new_argv)
                cmd = "source %s; cd %s/script; %s" % (GDSUtils.EXEC_ENV_FILE, install_dir, cmd0)
                result = {}
                cmd_with_ip = "export LD_LIBRARY_PATH=/lib64:$LD_LIBRARY_PATH; ssh %s \'%s\'" % (_ip, cmd)
                (status, output) = subprocess.getstatusoutput(cmd_with_ip)
                result[_ip] = output
                return result

            parall_result = GDSUtils.exec_command_remote_parallelly(ssh_cmd_handler,
                                                                    remote_hosts,
                                                                    timeout=GDSUtils.TIMEOUT)
            for key in parall_result:
                logger.log("[host %s]\n%s" % (key, parall_result[key]))
        except Exception as e:
            logger.logExit(str(e))
    if local_flag:
        # stop app locally
        if not is_process_exist(conf.port):
            logger.logExit("Find none process at port %d started by current user." % conf.port)
        # Remove watch dog firstly, then kill process.
        del_watch_dog(conf)
        cmd = "lsof -i :%d | grep -w '%s' | grep -v vi | grep -v grep| grep -v '/bin/bash' | grep -v '/bin/sh' | \
            awk '{print $2}' | xargs kill -9" % (conf.port, APP_NAME)
        (status, output) = subprocess.getstatusoutput(cmd)
        local_mark = "[Local]" if remote_hosts else ""
        logger.log(local_mark)
        if status != 0:
            logger.logExit("Failed to stop %s." % APP_NAME)
        logger.log("Stop %s at port %d [OK]" % (APP_NAME.upper(), conf.port))


def restart_gds(conf):
    """
    function: operation 'restart' handler, restart an app locally, dispatch to remote hosts with local mode.
              Done by gds_user.
    """
    stop_gds(conf)
    if local_flag:
        cmd = get_cmd_from_file(conf)
        (status, output) = subprocess.getstatusoutput(cmd)
        if status != 0:
            logger.logExit("Failed to restart %s at %s:%d." % (APP_NAME, conf.ip, conf.port))
        logger.log(output)


def clean_resource():
    """
    function: clean_resource, close files and so on.
    """
    try:
        logger.closeLog()
    except Exception:
        pass


if __name__ == '__main__':
    """
    main function
    """
    oper_dict = {"start": start_gds, "stop": stop_gds, "restart": restart_gds}
    # check if user is root
    if os.getuid() == 0:
        GaussLog.exitWithError(ErrorCode.GAUSS_501["GAUSS_50110"] % getpass.getuser())

    gds_user = getpass.getuser()
    try:
        parse_command_line()
        check_parameters()
        read_install_conf()
        oper_dict[curr_oper](curr_conf)
        clean_resource()
    except Exception as e:
        GaussLog.exitWithError(str(e))
