#!/usr/bin/env python3
# -*- coding=utf-8 -*-
# ############################################################################
# Copyright (c): 2012-2017, Huawei Tech. Co., Ltd.
# Description  : TaskPool.py is a utility to manage tasks.
# ############################################################################
try:
    import sys
    import os
    import signal
    import subprocess
    import threading
    import time

    sys.path.append(sys.path[0] + "/../../../")
    from gspylib.os.gsfile import g_file
except ImportError as ie:
    sys.exit("[GAUSS-52200] : Unable to import module: %s." % str(ie))

DIRECTORY_PERMISSION = 0o750


class WriterThread(threading.Thread):
    """
    class writer.
    Thread that processes the result content from TaskThread
     and writes the result content to a file.
    """

    def __init__(self, f_out, f_std):
        super(WriterThread, self).__init__()
        self.out_file = f_out
        self.err_file = f_std

        self.stdout = None
        self.stderr = None

    def run(self):
        """
        Writing the result content to a file.
        """
        if self.out_file:
            g_file.write_file_with_default_permission(self.out_file, self.stdout.encode(), False, "wb")

        if self.err_file:
            g_file.write_file_with_default_permission(self.err_file, self.stderr.encode(), False, "wb")


class TaskThread(threading.Thread):
    """
    class task
    Starts a task thread.
    """

    def __init__(self, host, cmd, f_out="", f_err="",
                 detail=False, timeout=0, shell_mode=False, inline=False):
        super(TaskThread, self).__init__()
        self.setDaemon(True)

        self.host = host
        self.cmd = cmd
        self.detail = bool(detail)
        self.timeout = timeout
        self.shell_mode = shell_mode
        self.inline = inline

        self.status = 0
        self.stdout, self.stderr = "", ""
        self.failures = []
        self.proc = None
        self.timestamp = time.time()
        self.isKill = False
        self.writer = WriterThread(f_out, f_err) if (f_out or f_err) else None

    def kill(self):
        """
        Kill the process of cmd.
        :param p: object of subprocess
        :return: NA
        """
        self.failures.append("Timed out")
        # kill process
        try:
            if self.proc:
                self.proc.kill()
        except (OSError, IOError):
            # If the kill fails, then just assume the process is dead.
            pass
        self.isKill = True
        # Set the status
        self.status = -1 * signal.SIGKILL
        self.failures.append("Killed by signal %s" % signal.SIGKILL)

    def getElapsedTime(self):
        """
         Getting elapsed timestamp.
        :return: timestamp
        """
        return time.time() - self.timestamp

    def checkTimeout(self):
        """
        check timed-out process.
        """
        if self.isKill or self.timeout <= 0:
            return False
        timeleft = self.timeout - self.getElapsedTime()
        if (timeleft <= 0):
            return True
        return False

    def run(self):
        """
        Execute the cmd on host.
        :return: NA
        """
        self.timestamp = time.time()
        env_dict = dict()
        for key in os.environ:
            env_dict[key] = os.environ.get(key)
        env_dict["LD_LIBRARY_PATH"] = "/lib64:%s" % os.environ.get("LD_LIBRARY_PATH")
        self.proc = subprocess.Popen(self.cmd, shell=False,
                                     stdout=subprocess.PIPE,
                                     stderr=subprocess.PIPE,
                                     env=env_dict)
        if self.timeout == 0:
            stdout, stderr = self.proc.communicate(timeout=None)
        else:
            stdout, stderr = self.proc.communicate(timeout=self.timeout)
        self.stdout += stdout.decode()
        self.stderr += stderr.decode()
        self.status = self.proc.returncode

    def __print_out(self):
        if not self.stdout and not self.stderr:
            return
        if self.shell_mode:
            sys.stderr.write("%s" % self.stderr)
            sys.stdout.write("%s" % self.stdout)
        else:
            if self.stdout:
                sys.stdout.write("%s: %s" % (self.host, self.stdout))
        # Use [-1] replace of .endswith, can avoid the problem about
        # coding inconsistencies
        if self.stdout and self.stdout[-1] != os.linesep:
            sys.stdout.write(os.linesep)
        if self.shell_mode and self.stderr and self.stderr[-1] != os.linesep:
            sys.stderr.write(os.linesep)

    def __print_result(self, index):
        """
        Print the result into sys.stdout
        :return: NA
        """
        if self.shell_mode:
            str_ = ""
        else:
            str_ = "[%s] %s [%s] %s" % (
                index,
                time.asctime().split()[3],
                "SUCCESS" if not self.status else "FAILURE",
                self.host
            )
            if self.status > 0:
                str_ += " Exited with error code %s" % self.status

        if self.failures:
            failures_msg = ", ".join(self.failures)
            str_ = str_ + " " + failures_msg

        if str_:
            print(str_)
        if self.inline:
            sys.stdout.write("%s" % self.stdout)

    def write(self, index):
        """
        Write the output into sys.stdout and files.
        :return: object of writer or None
        """
        # Print the stdout into sys.stdout
        if self.detail:
            self.__print_out()
        # Print the status
        self.__print_result(index)

        # Write the self.stdout and self.stderr into files.
        if self.writer:
            self.writer.stdout = self.stdout
            self.writer.stderr = self.stderr
            self.writer.start()
        return self.writer


class TaskPool(object):
    """
    class manager
    """

    def __init__(self, opts):
        """
        Initialize
        """
        self.out_path = opts.outdir
        self.err_path = opts.errdir
        self.detail = True
        self.parallel_num = opts.parallel
        self.timeout = opts.timeout
        self.shell_mode = opts.shellmode
        self.inline = opts.inline

        self.tasks = []
        self.running_tasks = []
        self.writers = []
        self.task_status = {}

    def __get_task_files(self, host):
        """
        Obtain the result file of the task.
        """
        std_path = ""
        if self.out_path:
            std_path = os.path.join(self.out_path, host)

        err_path = ""
        if self.err_path:
            err_path = os.path.join(self.err_path, host)

        return std_path, err_path

    def add_task(self, host, cmd):
        """
        Adding a Task to the Task Pool
        """

        f_out, f_err = self.__get_task_files(host)
        task = TaskThread(host, cmd, f_out, f_err, self.detail, self.timeout,
                          self.shell_mode, self.inline)
        self.tasks.append(task)

    def __get_writing_task(self):
        """
        Check the task status and obtain the running tasks.
        """
        still_running = []
        not_running = []
        time.sleep(0.1)
        # Check whether the task times out. If the task times out,
        # stop the task.
        for task in self.running_tasks:
            if task.checkTimeout():
                task.kill()

        # filter the still running tasks and not running tasks
        for task in self.running_tasks:
            if task.is_alive():
                still_running.append(task)
            else:
                self.task_status[task.host] = task.status
                not_running.append(task)

        # Start the writing thread of completed tasks
        for task in not_running:
            index = len(self.writers) + 1
            writer = task.write(index)
            if writer:
                self.writers.append(writer)

        self.running_tasks = still_running

    def __start_limit_task(self):
        """
        Starts the tasks within a specified number of parallel.
        """
        while self.tasks and len(self.running_tasks) < self.parallel_num:
            task = self.tasks.pop(0)
            self.running_tasks.append(task)
            task.start()

    def start(self):
        """
        Start to execute all tasks.
        """
        # Create the path of stdout and stderr
        if self.out_path and not os.path.exists(self.out_path):
            os.makedirs(self.out_path, DIRECTORY_PERMISSION)
        if self.err_path and not os.path.exists(self.err_path):
            os.makedirs(self.err_path, DIRECTORY_PERMISSION)

        # Do cmd
        while self.tasks or self.running_tasks:
            self.__get_writing_task()
            self.__start_limit_task()

        # Waiting for writing files complete.
        for writer in self.writers:
            writer.join()

        return self.task_status.values()


def read_host_file(host_file):
    """
    Reads the host file.
    Lines are of the form: host.
    Returns a list of host triples.
    """
    hosts = []
    try:
        with open(host_file) as fp:
            for line in fp:
                line = line.strip()
                if line or not line.startswith('#'):
                    hosts.append(line)
    except (OSError, IOError) as err:
        sys.stderr.write('Could not open hosts file: %s\n' % err)
        sys.exit(1)

    return hosts
