#
# Copyright (C) 2007-2016 CEA/DAM
# Copyright (C) 2015-2017 Stephane Thiell <sthiell@stanford.edu>
#
# This file is part of ClusterShell.
#
# ClusterShell is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# ClusterShell is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with ClusterShell; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

"""
ClusterShell worker interface.

A worker is a generic object which provides "grouped" work in a specific task.
"""

try:
    from inspect import getfullargspec  # py3
except ImportError:
    from inspect import getargspec as getfullargspec  # py2

import warnings

from ClusterShell.Worker.EngineClient import EngineClient
from ClusterShell.NodeSet import NodeSet
from ClusterShell.Engine.Engine import FANOUT_UNLIMITED, FANOUT_DEFAULT


def _eh_sigspec_invoke_compat(method, argc_legacy, *args):
    """
    Helper function to invoke an event handler method, with legacy
    signature compatibility if actual argc does match argc_legacy.
    This should be removed when old signatures (< 1.8) aren't supported
    anymore (in 2.x).
    """
    argc_actual = len(getfullargspec(method)[0])
    if argc_actual == argc_legacy:
        # Use legacy signature (1.x) deprecated as of 1.9
        warnings.warn("%s should use new %s() signature" % (method.__self__,
                                                            method.__name__),
                      DeprecationWarning)
        return method(*args[0:argc_legacy - 1])
    else:
        # Assume new signature (2.x)
        return method(*args)

def _eh_sigspec_ev_read_17(ev_read):
    """Helper function to check whether ev_read has the old 1.7 signature."""
    if len(getfullargspec(ev_read)[0]) == 2:
        warnings.warn("%s should use new ev_read() signature" % \
                      ev_read.__self__, DeprecationWarning)
        return True
    return False


class WorkerException(Exception):
    """Generic worker exception."""

class WorkerError(WorkerException):
    """Generic worker error."""

# DEPRECATED: WorkerBadArgumentError exception is deprecated as of 1.4,
# use ValueError instead.
WorkerBadArgumentError = ValueError

class Worker(object):
    """
    Worker is an essential base class for the ClusterShell library. The goal
    of a worker object is to execute a common work on a single or several
    targets (abstract notion) in parallel. Concrete targets and also the
    notion of local or distant targets are managed by Worker's subclasses
    (for example, see the DistantWorker base class).

    A configured Worker object is associated to a specific ClusterShell Task,
    which can be seen as a single-threaded Worker supervisor. Indeed, the work
    to be done is executed in parallel depending on other Workers and Task's
    current parameters, like current fanout value.

    ClusterShell is designed to write event-driven applications, and the Worker
    class is key here as Worker objects are passed as parameter of most event
    handlers (see the ClusterShell.Event.EventHandler class).

    Example of use:

        >>> from ClusterShell.Event import EventHandler
        >>> class MyOutputHandler(EventHandler):
        ...     def ev_read(self, worker, node, sname, msg):
        ...             print "%s: %s" % (node, line)
        ...
    """

    # The following common stream names are recognized by the Task class.
    # They can be changed per Worker, thus avoiding any Task buffering.
    SNAME_STDIN  = 'stdin'   #: stream name usually used for stdin
    SNAME_STDOUT = 'stdout'  #: stream name usually used for stdout
    SNAME_STDERR = 'stderr'  #: stream name usually used for stderr

    def __init__(self, handler):
        """Initializer. Should be called from derived classes."""
        # Associated EventHandler object
        self.eh = handler           #: associated :class:`.EventHandler`

        # Per Worker fanout value (positive integer).
        # Default is FANOUT_DEFAULT to use the fanout set at the Task level.
        # Change to FANOUT_UNLIMITED to always schedule this worker.
        # NOTE: the fanout value must be set before the Worker starts and
        # cannot currently be changed afterwards.
        self._fanout = FANOUT_DEFAULT

        # Update task rc? [private]
        # TODO: to be replaced with Task Event Handlers
        self._update_task_rc = True

        # Parent task (once bound)
        self.task = None            #: worker's task when scheduled or None
        self.started = False        #: set to True when worker has started
        self.metaworker = None
        self.metarefcnt = 0

        # current_x public variables (updated at each event accordingly)

        #: set to node in event handler; DEPRECATED: use :class:`.EventHandler` method argument **node**
        self.current_node = None
        #: set to stdout in event handler; DEPRECATED: use :class:`.EventHandler` method argument **msg** if ``sname==SNAME_STDOUT``
        self.current_msg = None
        #: set to stderr message in event handler; DEPRECATED: use :class:`.EventHandler` method argument **msg** if ``sname==SNAME_STDERR``
        self.current_errmsg = None
        #: set to return code in event handler; DEPRECATED: use :class:`.EventHandler` method argument **rc**
        self.current_rc = 0
        #: set to stream name in event handler; DEPRECATED: use :class:`.EventHandler` method argument **sname**
        self.current_sname = None

    def _set_task(self, task):
        """Bind worker to task. Called by task.schedule()."""
        if self.task is not None:
            # one-shot-only schedule supported for now
            raise WorkerError("worker has already been scheduled")
        self.task = task

    def _task_bound_check(self):
        """Helper method to check that worker is bound to a task."""
        if not self.task:
            raise WorkerError("worker is not task bound")

    def _engine_clients(self):
        """Return a list of underlying engine clients."""
        raise NotImplementedError("Derived classes must implement.")

    # Event generators

    def _on_start(self, key):
        """Called on command start."""
        self.current_node = key

        if not self.started:
            self.started = True
            if self.eh is not None:
                self.eh.ev_start(self)

        if self.eh is not None:
            _eh_sigspec_invoke_compat(self.eh.ev_pickup, 2, self, key)

    def _on_close(self, key, rc=None):
        """Called to generate events when the Worker is closing."""
        if self._update_task_rc:
            # rc may be None here for example when called from StreamClient
            # Only update task if rc is not None.
            if rc is not None:
                self.task._rc_set(self, key, rc)

        self.current_node = key
        self.current_rc = rc

        if self.eh is not None:
            _eh_sigspec_invoke_compat(self.eh.ev_hup, 2, self, key, rc)

    def _on_written(self, key, bytes_count, sname):
        """Notification of bytes written."""
        # set node and stream name (compat only)
        self.current_node = key
        self.current_sname = sname

        if self.eh is not None:
            self.eh.ev_written(self, key, sname, bytes_count)

    # Base getters

    def did_timeout(self):
        """Return whether this worker has aborted due to timeout."""
        self._task_bound_check()
        return self.task._num_timeout_by_worker(self) > 0

    def read(self, node=None, sname='stdout'):
        """Read worker stream buffer.

        Return stream read buffer of current worker.

        Arguments:
        :param node: node name, can also be set to None for simple worker
        having worker.key defined (default is None)
        :param sname: stream name (default is 'stdout')
        """
        self._task_bound_check()
        return self.task._msg_by_source(self, node, sname)

    # Base actions

    def abort(self):
        """Abort processing any action by this worker.

        Safe to call on an already closing or aborting worker.
        """
        raise NotImplementedError("Derived classes must implement.")

    def flush_buffers(self):
        """Flush any messages associated to this worker."""
        self._task_bound_check()
        self.task._flush_buffers_by_worker(self)

    def flush_errors(self):
        """Flush any error messages associated to this worker."""
        self._task_bound_check()
        self.task._flush_errors_by_worker(self)


class DistantWorker(Worker):
    """Base class DistantWorker.

    DistantWorker provides a useful set of setters/getters to use with
    distant workers like ssh or pdsh.
    """

    # Event generators

    def _on_node_msgline(self, node, msg, sname):
        """Message received from node, update last* stuffs."""
        # Maxoptimize this method as it might be called very often.
        task = self.task
        assert not isinstance(node, NodeSet) # for testing
        # update task msgtree
        task._msg_add(self, node, sname, msg)

        ### LEGACY (1.x) ###
        # set stream name
        self.current_sname = sname

        # generate event
        self.current_node = node
        if sname == self.SNAME_STDERR:
            self.current_errmsg = msg
            if self.eh is not None:
                # call old ev_error for compat (default is no-op)
                if hasattr(self.eh, 'ev_error'):  # missing in 1.8.0!
                    self.eh.ev_error(self)
                # /!\ NOT elif
                if not _eh_sigspec_ev_read_17(self.eh.ev_read):
                    ### FUTURE (2.x) ###
                    self.eh.ev_read(self, node, sname, msg)
        else:
            self.current_msg = msg
            if self.eh is not None:
                # ev_read: check for old signature first (< 1.8)
                if _eh_sigspec_ev_read_17(self.eh.ev_read):
                    self.eh.ev_read(self)
                else:
                    ### FUTURE (2.x) ###
                    self.eh.ev_read(self, node, sname, msg)

    def _on_node_close(self, node, rc):
        """Command return code received."""
        Worker._on_close(self, node, rc)

    def _on_node_timeout(self, node):
        """Update on node timeout."""
        # Update current_node to allow node resolution after ev_timeout.
        self.current_node = node

        self.task._timeout_add(self, node)

    def node_buffer(self, node):
        """Get specific node buffer."""
        return self.read(node, self.SNAME_STDOUT)

    def node_error(self, node):
        """Get specific node error buffer."""
        return self.read(node, self.SNAME_STDERR)

    node_error_buffer = node_error

    def node_retcode(self, node):
        """
        Get specific node return code.

        :raises KeyError: command on node has not yet finished (no return code
            available), or this node is not known by this worker
        """
        self._task_bound_check()
        try:
            rc = self.task._rc_by_source(self, node)
        except KeyError:
            raise KeyError(node)
        return rc

    node_rc = node_retcode

    def iter_buffers(self, match_keys=None):
        """
        Returns an iterator over available buffers and associated
        NodeSet. If the optional parameter match_keys is defined, only
        keys found in match_keys are returned.
        """
        self._task_bound_check()
        for msg, keys in self.task._call_tree_matcher(
                self.task._msgtree(self.SNAME_STDOUT).walk, match_keys, self):
            yield msg, NodeSet.fromlist(keys)

    def iter_errors(self, match_keys=None):
        """
        Returns an iterator over available error buffers and associated
        NodeSet. If the optional parameter match_keys is defined, only
        keys found in match_keys are returned.
        """
        self._task_bound_check()
        for msg, keys in self.task._call_tree_matcher(
                self.task._msgtree(self.SNAME_STDERR).walk, match_keys, self):
            yield msg, NodeSet.fromlist(keys)

    def iter_node_buffers(self, match_keys=None):
        """
        Returns an iterator over each node and associated buffer.
        """
        self._task_bound_check()
        return self.task._call_tree_matcher(
            self.task._msgtree(self.SNAME_STDOUT).items, match_keys, self)

    def iter_node_errors(self, match_keys=None):
        """
        Returns an iterator over each node and associated error buffer.
        """
        self._task_bound_check()
        return self.task._call_tree_matcher(
            self.task._msgtree(self.SNAME_STDERR).items, match_keys, self)

    def iter_retcodes(self, match_keys=None):
        """
        Returns an iterator over return codes and associated NodeSet.
        If the optional parameter match_keys is defined, only keys
        found in match_keys are returned.
        """
        self._task_bound_check()
        for rc, keys in self.task._rc_iter_by_worker(self, match_keys):
            yield rc, NodeSet.fromlist(keys)

    def iter_node_retcodes(self):
        """
        Returns an iterator over each node and associated return code.
        """
        self._task_bound_check()
        return self.task._krc_iter_by_worker(self)

    def num_timeout(self):
        """
        Return the number of timed out "keys" (ie. nodes) for this worker.
        """
        self._task_bound_check()
        return self.task._num_timeout_by_worker(self)

    def iter_keys_timeout(self):
        """
        Iterate over timed out keys (ie. nodes) for a specific worker.
        """
        self._task_bound_check()
        return self.task._iter_keys_timeout_by_worker(self)

class StreamClient(EngineClient):
    """StreamWorker's default EngineClient.

    StreamClient is the EngineClient subclass used by default by
    StreamWorker. It handles some generic methods to pass data to the
    StreamWorker.
    """

    def _start(self):
        """Called on EngineClient start."""
        assert not self.worker.started
        self.worker._on_start(self.key)
        return self

    def _read(self, sname, size=65536):
        """Read data from process."""
        return EngineClient._read(self, sname, size)

    def _close(self, abort, timeout):
        """Close client. See EngineClient._close()."""
        EngineClient._close(self, abort, timeout)
        if timeout:
            assert abort, "abort flag not set on timeout"
            self.worker._on_timeout(self.key)

        # return code not available
        self.worker._on_close(self.key)

        if self.worker.eh:
            _eh_sigspec_invoke_compat(self.worker.eh.ev_close, 2, self, timeout)

    def _handle_read(self, sname):
        """Engine is telling us there is data available for reading."""
        # Local variables optimization
        task = self.worker.task
        msgline = self.worker._on_msgline

        debug = task.info("debug", False)
        if debug:
            print_debug = task.info("print_debug")
            for msg in self._readlines(sname):
                print_debug(task, "LINE %s" % msg)
                msgline(self.key, msg, sname)
        else:
            for msg in self._readlines(sname):
                msgline(self.key, msg, sname)

    def _flush_read(self, sname):
        """Called at close time to flush stream read buffer."""
        stream = self.streams[sname]
        if stream.readable() and stream.rbuf:
            # We still have some read data available in buffer, but no
            # EOL. Generate a final message before closing.
            self.worker._on_msgline(self.key, stream.rbuf, sname)

    def write(self, buf, sname=None):
        """Write to writable stream(s)."""
        if sname is not None:
            self._write(sname, buf)
            return
        # sname not specified: "broadcast" to all writable streams...
        for writer in self.streams.writers():
            self._write(writer.name, buf)

    def set_write_eof(self, sname=None):
        """Set EOF flag to writable stream(s)."""
        if sname is not None:
            self._set_write_eof(sname)
            return
        # sname not specified: set eof flag on all writable streams...
        for writer in self.streams.writers():
            self._set_write_eof(writer.name)

class StreamWorker(Worker):
    """StreamWorker base class [v1.7+]

    The StreamWorker class implements a base (but concrete) Worker that
    can read and write to multiple streams. Unlike most other Workers,
    it does not execute any external commands by itself. Rather, it
    should be pre-bound to "streams", ie. file(s) or file descriptor(s),
    using the two following methods:

        >>> worker.set_reader('stream1', fd1)
        >>> worker.set_writer('stream2', fd2)

    Like other Workers, the StreamWorker instance should be associated
    with a Task using task.schedule(worker). When the task engine is
    ready to process the StreamWorker, all of its streams are being
    processed together. For that reason, it is not possible to add new
    readers or writers to a running StreamWorker (ie. task is running
    and worker is already scheduled).

    Configured readers will generate ev_read() events when data is
    available for reading, with the stream name passed as one of its
    argument.

    Configured writers will allow the use of the method write(), eg.
    worker.write(data, 'stream2'), to write to the stream.
    """

    def __init__(self, handler, key=None, stderr=False, timeout=-1,
                 autoclose=False, client_class=StreamClient):
        Worker.__init__(self, handler)
        if key is None: # allow key=0
            key = self
        self.clients = [client_class(self, key, stderr, timeout, autoclose)]

    def set_reader(self, sname, sfile, retain=True, closefd=True):
        """Add a readable stream to StreamWorker.

        Arguments:
        :param sname: the name of the stream (string)
        :param sfile: the stream file or file descriptor
        :param retain: whether the stream retains engine client
        (default is True)
        :param closefd: whether to close fd when the stream is closed
        (default is True)
        """
        if not self.clients[0].registered:
            self.clients[0].streams.set_reader(sname, sfile, retain, closefd)
        else:
            raise WorkerError("cannot add new stream at runtime")

    def set_writer(self, sname, sfile, retain=True, closefd=True):
        """Set a writable stream to StreamWorker.

        Arguments:
        :param sname: the name of the stream (string)
        :param sfile: the stream file or file descriptor
        :param retain: whether the stream retains engine client
        (default is True)
        :param closefd: whether to close fd when the stream is closed
        (default is True)
        """
        if not self.clients[0].registered:
            self.clients[0].streams.set_writer(sname, sfile, retain, closefd)
        else:
            raise WorkerError("cannot add new stream at runtime")

    def _engine_clients(self):
        """Return a list of underlying engine clients."""
        return self.clients

    def set_key(self, key):
        """Source key for this worker is free for use.

        Use this method to set the custom source key for this worker.
        """
        self.clients[0].key = key

    def _on_msgline(self, key, msg, sname):
        """Add a message."""
        # update task msgtree
        self.task._msg_add(self, key, sname, msg)

        ### LEGACY (1.x) ###
        # set stream name
        self.current_sname = sname

        # generate event
        if sname == 'stderr':
            self.current_errmsg = msg
            if self.eh is not None:
                # call old ev_error for compat (default is no-op)
                if hasattr(self.eh, 'ev_error'):  # missing in 1.8.0!
                    self.eh.ev_error(self)
                # /!\ NOT elif
                if not _eh_sigspec_ev_read_17(self.eh.ev_read):
                    ### FUTURE (2.x) ###
                    self.eh.ev_read(self, key, sname, msg)
        else:
            self.current_msg = msg
            if self.eh is not None:
                # ev_read: check for old signature first (< 1.8)
                if _eh_sigspec_ev_read_17(self.eh.ev_read):
                    self.eh.ev_read(self)
                else:
                    ### FUTURE (2.x) ###
                    self.eh.ev_read(self, key, sname, msg)

    def _on_timeout(self, key):
        """Update on timeout."""
        self.task._timeout_add(self, key)

        # trigger timeout event (deprecated in 1.8+)
        # also use hasattr check because ev_timeout was missing in 1.8.0
        if self.eh and hasattr(self.eh, 'ev_timeout'):
            warnings.warn("%s should use new ev_close() instead of " \
                          "ev_timeout()" % self.eh, DeprecationWarning)
            self.eh.ev_timeout(self)

    def abort(self):
        """Abort processing any action by this worker.

        Safe to call on an already closing or aborting worker.
        """
        self.clients[0].abort()

    def read(self, node=None, sname='stdout'):
        """Read worker stream buffer.

        Return stream read buffer of current worker.

        Arguments:
        :param node: node name, can also be set to None for simple worker
        having worker.key defined (default is None)
        :param sname: stream name (default is 'stdout')
        """
        return Worker.read(self, node or self.clients[0].key, sname)

    def write(self, buf, sname=None):
        """Write to worker.

        If sname is specified, write to the associated stream,
        otherwise write to all writable streams.
        """
        self.clients[0].write(buf, sname)

    def set_write_eof(self, sname=None):
        """
        Tell worker to close its writer file descriptor once flushed.

        Do not perform writes after this call. Like write(), sname can
        be optionally specified to target a specific writable stream,
        otherwise all writable streams are marked as EOF.
        """
        self.clients[0].set_write_eof(sname)

class WorkerSimple(StreamWorker):
    """WorkerSimple base class [DEPRECATED]

    Implements a simple Worker to manage common process
    stdin/stdout/stderr streams.

    [DEPRECATED] use StreamWorker.
    """

    def __init__(self, file_reader, file_writer, file_error, key, handler,
                 stderr=False, timeout=-1, autoclose=False, closefd=True,
                 client_class=StreamClient):
        """Initialize WorkerSimple worker."""
        StreamWorker.__init__(self, handler, key, stderr, timeout, autoclose,
                              client_class=client_class)
        if file_reader:
            self.set_reader('stdout', file_reader, closefd=closefd)
        if file_error:
            self.set_reader('stderr', file_error, closefd=closefd)
        if file_writer:
            self.set_writer('stdin', file_writer, closefd=closefd)
        # keep reference of provided file objects during worker lifetime
        self._filerefs = (file_reader, file_writer, file_error)

    def error_fileno(self):
        """Return the standard error reader file descriptor (integer)."""
        return self.clients[0].streams['stderr'].fd

    def reader_fileno(self):
        """Return the reader file descriptor (integer)."""
        return self.clients[0].streams['stdout'].fd

    def writer_fileno(self):
        """Return the writer file descriptor as an integer."""
        return self.clients[0].streams['stdin'].fd

    def error(self):
        """Read worker error buffer."""
        return self.read(sname='stderr')

    def _on_start(self, key):
        """Called on command start."""
        if not self.started:
            self.started = True
            if self.eh is not None:
                self.eh.ev_start(self)

        if self.eh is not None:
            # generate ev_pickup
            _eh_sigspec_invoke_compat(self.eh.ev_pickup, 2, self, key)

    def _on_close(self, key, rc=None):
        """Called to generate events when the Worker is closing."""
        self.current_rc = rc

        # rc may be None here for example when called from StreamClient
        # Only update task if rc is not None.
        if rc is not None:
            self.task._rc_set(self, key, rc)

        if self.eh is not None:
            # generate ev_hup
            _eh_sigspec_invoke_compat(self.eh.ev_hup, 2, self, key, rc)

    def _on_written(self, key, bytes_count, sname):
        """Notification of bytes written."""
        # set node and stream name (compat only)
        self.current_sname = sname

        if self.eh is not None:
            # generate ev_written
            self.eh.ev_written(self, key, sname, bytes_count)