from __future__ import absolute_import
from twisted.internet.protocol import ProcessProtocol
from twisted.protocols.basic import LineReceiver
from .twisted import TwistedTask
from ..sparts import counter, option

import functools
import signal
import time
import twisted.python.threadable
import twisted.internet.threads


class CommandTask(TwistedTask):
    LOOPLESS = True
    OPT_PREFIX = 'cmd'

    kill_timeout = option(type=float, default=10.0,
                          help="Default shutdown kill timeout for outstanding "
                               "commands [%(default)s]")
    started = counter()
    finished = counter()

    def initTask(self):
        super(CommandTask, self).initTask()
        self.outstanding = {}

    def procExited(self, on_exit, proto, trans, reason):
        self.logger.debug("%s closed for %s", trans, reason)
        if on_exit is not None:
            on_exit(reason)

        self.outstanding.pop(trans)

        self.finished.increment()
        return None

    def run(self, command, on_stdout=None, on_stderr=None, on_exit=None,
            line_buffered=True, kill_timeout=None):
        self.logger.debug("task starting %s...", command)
        if isinstance(command, basestring):
            command = [command]

        # wrap on_exit with helper to remove registered comments
        on_exit = functools.partial(self.procExited, on_exit)

        proto = _ProcessProtocolAdapter(on_stdout, on_stderr, on_exit,
                                       line_buffered=line_buffered)

        if twisted.python.threadable.isInIOThread():
            result = self.reactor.spawnProcess(proto, executable=command[0],
                                               args=command)
        else:
            result = twisted.internet.threads.blockingCallFromThread(
                self.reactor, self.reactor.spawnProcess, proto,
                executable=command[0], args=command)

        self.outstanding[result] = kill_timeout

        self.started.increment()
        return result

    def join(self):
        """Override to wait for process workers to shutdown / be killed"""
        # TODO: Conditions instead of sleep polling? 
        while len(self.outstanding) > 0:
            time.sleep(0.250)

    def _killOutstanding(self, trans):
        if trans in self.outstanding:
            self.logger.info("Sending SIGKILL to %s", trans)
            trans.signalProcess(signal.SIGKILL)

    def stop(self):
        # twisted is pretty smart; the default signal handlers it installs
        # propagate SIGTERM to its children, so while we don't need to manually
        # TERM, we might still need to set some kill timeouts
        super(CommandTask, self).stop()

        for trans, kill_timeout in self.outstanding.iteritems():
            if kill_timeout is None:
                kill_timeout = self.kill_timeout

            self.logger.info("Enqueuing kill for %s in %.1fs",
                             trans, kill_timeout)
            args = (kill_timeout, self._killOutstanding, trans)
            if twisted.python.threadable.isInIOThread():
                self.reactor.callLater(*args)
            else:
                self.reactor.callFromThread(self.reactor.callLater, *args)


class _ProcessProtocolAdapter(ProcessProtocol):
    def __init__(self, on_stdout=None, on_stderr=None, on_exit=None,
                 line_buffered=True):

        self.line_buffered = line_buffered
        self.on_stdout = on_stdout or (lambda *args: None)
        self.on_stderr = on_stderr or (lambda *args: None)
        self.on_exit = on_exit or (lambda *args: None)

        if self.line_buffered:
            self.stdout_buffer = self._makeBuffer(self.on_stdout)
            self.stderr_buffer = self._makeBuffer(self.on_stderr)

    def _makeBuffer(self, callback):
        buf = LineReceiver()
        buf.delimiter = "\n"
        buf.lineReceived = lambda line: callback(self.transport, line)
        return buf

    def connectionMade(self):
        ProcessProtocol.connectionMade(self)
        self.original_pid = self.transport.pid

    def outReceived(self, data):
        if self.line_buffered:
            self.stdout_buffer.dataReceived(data)
        else:
            self.on_stdout(self.transport, data)

    def errReceived(self, data):
        if self.line_buffered:
            self.stderr_buffer.dataReceived(data)
        else:
            self.on_stderr(self.transport, data)

    def processEnded(self, reason):
        # Ugh, this is a hack... I need to pass self because Twisted unsets pid
        # after the process has exited, and I need the *old* pid to do proper
        # accounting.  *YES* I get why this is fundamentally broken
        self.on_exit(self, self.transport, reason)
        self.transport.loseConnection()

@@ -2765,20 +2765,16 @@
 tanding.
-iter
 items():

Fix race between cmdtask shutdown and process closure handler