from __future__ import print_function
import socket
import os
import sys
import subprocess
import signal
import re
import shlex

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET

# Exit Code
EXIT_SUCCESS = 0
EXIT_FAILURE = -1
EXIT_STATUS_IDEL = 1
EXIT_STATUS_RUNNING = 2

POPEN_TIMEOUT = 60

g_cfg_path = ""
g_bin_path = ""


class gds_entry:
    port = -1
    name = None
    log_file = None
    data_dir = None
    err_dir = None
    data_seg = None
    err_seg = None
    ctl_file = None
    recursive = False
    daemon = False
    host = None
    ip = None
    parallel = -1
    help = False

    def is_valid(self):
        if self.help is True:
            return True
        if self.port < 1024 or self.port > 65535:
            return False
        if self.data_dir is None:
            return False
        if self.name is None:
            return False
        if re.search('^\w+([-+.]\w+)*$', self.name) is None:
            return False

        return True

    def to_str(self):
        currStr = ('name:{0}'.format(self.name))
        if self.ip is not None:
            currStr += (' ip:{0}'.format(self.ip))

        currStr += (' port:{0:d}'.format(self.port))
        currStr += (' data_dir:{0}'.format(self.dta_dir))
        if self.err_dir is not None:
            currStr += (' err_dir:{0}'.format(self.drr_dir))

        if self.data_seg is not None:
            currStr += (' data_seg:{0}'.format(self.data_seg))

        if self.err_seg is not None:
            currStr += (' err_seg:{0}'.format(self.err_seg))

        if self.log_file is not None:
            currStr += (' log_file:{0}'.format(self.log_file))

        if self.host is not None:
            currStr += (' host:{0}'.format(self.host))

        if self.recursive is True:
            currStr += ' recursive:{true}'
        else:
            currStr += ' recursive:{false}'

        if self.daemon is True:
            currStr += ' daemon:{true}'
        else:
            currStr += ' daemon:{false}'

        if self.enablessl is True:
            currStr += ' enablessl:{true}'
        else:
            currStr += ' enablessl:{false}'

        if self.parallel is not None:
            currStr += (' parallel:{0}'.format(self.parallel))
        if self.help is True:
            currStr += ' help:{true}'
        else:
            currStr += ' help:{false}'

        return currStr


def usage():
    print("gds_ctl.py [start|stop|status]")


# Read the GDS configration
def read_conf(filename):
    try:
        tree = ET.ElementTree(file=filename)
        root = tree.getroot()
        gds_array = []
        gds_dict = {}

    except Exception:
        print('Invalid gds configuration file \"' + filename + '\"')
        sys.exit(EXIT_FAILURE)
    for cell in root:

        entry = gds_entry()
        entry.port = int(cell.get('port', '-1'))
        entry.data_dir = cell.get('data_dir', None)
        entry.err_dir = cell.get('err_dir', None)
        entry.data_seg = cell.get('data_seg', None)
        entry.err_seg = cell.get('err_seg', None)
        entry.recursive = (cell.get('recursive', 'False').lower() == 'true')
        entry.daemon = (cell.get('daemon', 'False').lower() == 'true')
        entry.help = (cell.get('help', 'False').lower() == 'true')
        # make the data_dir as a ABS path
        if entry.data_dir is not None:
            entry.data_dir = entry.data_dir.strip()
            if os.path.isabs(entry.data_dir) is False:
                entry.data_dir = sys.path[0] + '/' + entry.data_dir

        # make the err_dir as a ABS path
        if entry.err_dir is not None:
            entry.err_dir = entry.err_dir.strip()
            if os.path.isabs(entry.err_dir) is False:
                entry.err_dir = sys.path[0] + '/' + entry.err_dir

        entry.name = cell.get('name', None)
        entry.ctl_file = ".gds.%s" % entry.name
        entry.log_file = cell.get('log_file', None)
        entry.host = cell.get('host', None)
        entry.ip = cell.get('ip', None)
        entry.parallel = cell.get('parallel', None)

        if entry.is_valid() is not True:
            print('Invalid item in configuration file \"' + filename + '\"')
            sys.exit(EXIT_FAILURE);
        gds_array.append(entry)
        if entry.name in gds_dict.keys():
            print('Conflict GDS name \"{0}\"'.format(entry.name))
            sys.exit(EXIT_FAILURE)
        gds_dict[entry.name] = entry

    gds_dict.clear()
    return gds_array


# Start GDS
def gds_start(entries):
    for cell in entries:
        if cell.help is True:
            cmd = g_bin_path + '/gds -h'
            if os.system(cmd) != 0:
                return EXIT_FAILURE
            else:
                return EXIT_SUCCESS

        if __is_GDS_exist__(cell):
            print(('GDS {0} already exists.').format(cell.name))
            return EXIT_FAILURE
        # Create the data directory if dosen't exist,
        if os.path.exists(cell.data_dir) is False:
            os.makedirs(cell.data_dir)

        # Create the error log directory if dosen't exist,
        if os.path.exists(cell.err_dir) is False:
            os.makedirs(cell.err_dir)

        if os.path.exists(cell.ctl_file):
            os.system(('rm {0}'.format(cell.ctl_file)))

        cmd = g_bin_path + ('/gds -d {0} -s {1} -D'.format(cell.data_dir, cell.ctl_file))
        if cell.ip is not None:
            cmd += (' -p {0}:{1:d}'.format(cell.ip, cell.port))
        else:
            cmd += (' -p {0:d}'.format(cell.port))
        if cell.log_file is not None:
            cmd += (' -l {0}'.format(cell.log_file))
        if cell.recursive is True:
            cmd += ' -r'
        if cell.daemon is True:
            cmd += ' -D'
        if cell.host is not None:
            cmd += (' -H {0}'.format(cell.host))

        if cell.parallel != -1:
            cmd += (' -t {0}'.format(cell.parallel))

        if cell.err_dir is not None:
            cmd += (' -e {0}'.format(cell.err_dir))

        if cell.data_seg is not None:
            cmd += (' -S {0}'.format(cell.data_seg))

        if cell.err_seg is not None:
            cmd += (' -E {0}'.format(cell.err_seg))

        if os.system(cmd) != 0:
            print('Start GDS {0:20s}[ERROR]'.format(cell.name))
            return EXIT_FAILURE
        print('Start GDS {0:30s}[OK]'.format(cell.name))
    return EXIT_SUCCESS


# Stop the GDS
def gds_stop_entries(entries):
    gds_proc = []
    cmd = 'ps ux'
    out, err = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE).communicate(timeout=POPEN_TIMEOUT)
    for line in out.decode().splitlines():
        if 'gds' in line:
            gds_proc.append(line)

    for cell in entries:
        for proc in gds_proc:
            if cell.ctl_file in proc.split(' '):
                pid = int(proc.split(None, 2)[1])
                os.kill(pid, signal.SIGRTMIN)
                if (os.path.exists(cell.ctl_file)):
                    os.system(('rm {0}'.format(cell.ctl_file)))
                if __is_GDS_exist__(cell) is False:
                    print('Stop GDS {0:32s}[OK]'.format(cell.name))
                else:
                    print('Stop GDS {0:32s}[ERROR]'.format(cell.name))
    return EXIT_SUCCESS


def gds_stop_all():
    cmd = 'ps ux'
    out, err = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE).communicate(timeout=POPEN_TIMEOUT)

    for line in out.decode().splitlines():
        if 'gds' in line and '-p' in line and '-d' in line:
            process_info = line.split(' ')
            i = 0
            for p in process_info:
                if p == '-p':
                    break
                i = i + 1
            single = process_info[i + 1]
            pid = int(line.split(None, 2)[1])
            os.kill(pid, signal.SIGRTMIN)
            if __is_GDS_exist__(pid) is False:
                print('Stop GDS {0:32s}[OK]'.format(single))
            else:
                print('Stop GDS {0:32s}[ERROR]'.format(single))
    return EXIT_SUCCESS


def gds_stop_single(single):
    cmd = 'ps ux'
    out, err = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE).communicate(timeout=POPEN_TIMEOUT)
    for line in out.decode().splitlines():
        if 'gds' in line and '-p' in line and '-d' in line:
            process_info = line.split(' ')
            i = 0
            for p in process_info:
                if p == '-p':
                    break
                i = i + 1
            if process_info[i + 1] != single:
                continue
            pid = int(line.split(None, 2)[1])
            os.kill(pid, signal.SIGRTMIN)
            if __is_GDS_exist__(pid) is False:
                print('Stop GDS {0:32s}[OK]'.format(single))
            else:
                print('Stop GDS {0:32s}[ERROR]'.format(single))
    return EXIT_SUCCESS


def __is_GDS_exist__(entry):
    cmd = 'ps ux'
    out, err = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE).communicate(timeout=POPEN_TIMEOUT)
    for line in out.decode().splitlines():
        if ('gds' in line) and (entry.ctl_file in line.split(' ')):
            return True
    return False


def __is_GDS_exist__(pid):
    cmd = "ps ux"
    out, err = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE).communicate(timeout=POPEN_TIMEOUT)
    for line in out.decode().splitlines():
        if ('gds' in line) and (pid in line.split(' ')):
            return True
    return False


# Notice: Now the 'STATUS' command can only return the status of the first GDS in configuration file.
def gds_status(self):
    gds = self[0]
    if gds is None:
        return EXIT_FAILURE;
    if os.path.exists(gds.ctl_file) is False:
        return EXIT_FAILURE
    try:
        client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        client.settimeout(5)
        client.connect(gds.ctl_file)
        client.send("Q".encode())
        result = client.recv(1).decode()
    except socket.timeout:
        client.close()
        return EXIT_FAILURE
    except socket.error:
        return EXIT_FAILURE
    if result == 'I':
        print(('GDS {0:32}[IDEL]').format(gds.name))
        return EXIT_STATUS_IDEL
    elif result == 'R':
        print(('GDS {0:32}[RUNNING]').format(gds.name))
        return EXIT_STATUS_RUNNING
    else:
        return EXIT_FAILURE


def main():
    global g_cfg_path
    global g_bin_path

    g_bin_path = sys.path[0]
    # The configuration file gds.cfg is under '$install_dir/config/' by default.
    g_cfg_paths = [
        os.path.join(sys.path[0], './../config/gds.conf'),  # since 8.1.0
        os.path.join(sys.path[0], './config/gds.conf'),     # less than 8.1.0, such as 8.0
        os.path.join(sys.path[0], 'gds.conf')               # other scenarios
    ]
    for path in g_cfg_paths:
        if os.path.exists(path):
            g_cfg_path = path
            break

    # Verify if the configuration file exists
    if g_cfg_path == '':
        print(f"Can't find GDS configure file from candidate paths: {g_cfg_paths}")
        sys.exit(EXIT_FAILURE)
    if len(sys.argv) == 1:
        usage()
        sys.exit(EXIT_FAILURE)
    cmd = sys.argv[1]
    if str(cmd).lower() == 'start':
        sys.exit(gds_start(read_conf(g_cfg_path)))
    elif str(cmd).lower() == 'stop':
        # stop gds entries
        if len(sys.argv) == 2:
            sys.exit(gds_stop_entries(read_conf(g_cfg_path)))
        if len(sys.argv) == 3:
            cmd = sys.argv[2]
            # stop all gds which could be stopped
            if cmd == "all":
                sys.exit(gds_stop_all())
            # stop single gds specified by "[ip]:port"
            else:
                sys.exit(gds_stop_single(cmd))
    elif str(cmd).lower() == 'status':
        sys.exit(gds_status(read_conf(g_cfg_path)))
    else:
        print("Invalid commands \"" + cmd + "\"")
        usage()
        sys.exit(EXIT_FAILURE)


if __name__ == "__main__":
    main()
