"""Spawn subprocesses via a pty.

   The normal way to use this is by creating an instance of the Spawn
   class.  This will connect the stdin, stdout and stderr of the new
   process to the pty.  You can then use the send() method to send
   data to the procces, and wait for input using the normal Expectable
   API.

   If you need to differentiate between stdout and stderr of the
   proces, you can use the spawn2() function.  This will create a
   pipe and a pty.  The pty will be connected to stdin and stdout of
   the process, while the pipe will be connected to stderr.  spawn2()
   returns two expectable objects, representing the pty and the pipe.

   The method of connecting stderr to a pipe is not perfect.  Some
   programs assumes that stderr is connected to the tty, and will not
   work properly if it is connected to a pipe.  For instance, if you
   start csh on Solaris 9 this way, it will not have job control
   features.  And if you start bash, and try to run stty from bash, it
   will fail (tested with bash-2.05 under Solaris 9 and Linux 2.4.22).
   As long as the program you start isn't a shell or a program that
   calls isatty(2) you should be fine, though.

   inherit_stty
   ============

   By default, the new pty will inherit the stty settings of stdin, if
   stdin is a tty.  You can disable this by setting the global module
   variable inherit_stty to False.  The setting of inherit_stty can
   also be overriden by the optional argument inherit_stty to the
   Spawn constructor or spawn2() function.

   stty_init
   =========

   After possibly inheriting the stty settings (see inherit_stty
   above) stty will be run on the new pty with the stty_init string as
   an argument.  stty_init is default False, which means that this
   step is skipped.  stty_init is also a global variable that can be
   overridden by optional arguments to the Spawn constructor or
   spawn2() function.

"""

import os
import sys
import pty
import fcntl

import pcl_expect
import pcl_expect.remote_exception

__all__ = [
    "inherit_stty",
    "stty_init",
    "Spawn",
    "spawn2",
    ]

inherit_stty = True
stty_init = False

def set_cloexec_flag(fd):
    """Set the FD_CLOEXEC flag on fd."""
    old = fcntl.fcntl(fd, fcntl.F_GETFD)
    fcntl.fcntl(fd, fcntl.F_SETFD, old | fcntl.FD_CLOEXEC)

def _use(a, b):
    """Return b if a is None, otherwise return a."""
    if a is None:
        return b
    else:
        return a

def _inherit_stty(override):
    """Return the inherit_stty setting to use."""
    _use(override, inherit_stty)

def _stty_init(override):
    """Return the stty_init setting to use."""
    _use(override, stty_init)


def _spawn(cmd, inherit_stty, stty_init, use_stderr_pipe):
    """Spawn a process running cmd.

       cmd -- A list of strings to be used as the argv vector of
              os.execvp.  cmd[0] will also be used as the program to
              run.  This can also be a string that will be split on
              space to produce this list.  There is no way to quote
              a space, so unless the string is a constant you should
              normally not use that form.

       inherit_stty -- True if the stty settings should be inherited
              if stdin is a tty.

       stty_init -- String to pass to stty to initialize the pty.

       use_stderr_pipe -- True if stderr of the new process should be
              connected to a separate pipe instead of to the pty.

    """

    inherit = inherit_stty and os.isatty(sys.stdin.fileno())
    if inherit:
        f = os.popen("stty -g")
        settings = f.read()
        f.close()

    if isinstance(cmd, basestring):
        cmd = cmd.split(" ")

    # Create a synchronization pipe, so that the parent doesn't
    # continue execution until the pty has been initialized and the
    # new process execed.  If all goes well, the parent will receive
    # an end-of-file on the pipe when the exec closes the pipe.  If
    # exec fails, the errno code will be sent via this pipe.
    (r, w) = os.pipe()
    set_cloexec_flag(w)

    if use_stderr_pipe:
        (stderr_r, stderr_w) = os.pipe()
    else:
        stderr_r = None

    pid, fd = pty.fork()
    if pid == 0:
        try:
            try:
                os.close(r)
                if inherit:
                    os.system("stty %s" % settings)
                if stty_init is not None:
                    os.system("stty %s" % stty_init)
                if use_stderr_pipe:
                    os.close(stderr_r)
                    os.dup2(stderr_w, 2)
                    os.close(stderr_w)

                os.execvp(cmd[0], cmd)
            except:
                os.write(w, remote_exception.serialize())
        finally:
            os._exit(1)

    # In parent.
    os.close(w)
    if use_stderr_pipe:
        os.close(stderr_w)
    res = os.fdopen(r).read()
    if res != "":
        os.close(fd)
        os.waitpid(pid, 0)
        remote_exception.re_raise(res)
    return pid, fd, stderr_r

class Spawn(pcl_expect.Expectable):
    def __init__(self, cmd, inherit_stty = None, stty_init = None,
                 use_stderr_pipe = False):

        self.__child_pid, fd, stderr = _spawn(
            cmd,
            _inherit_stty(inherit_stty),
            _stty_init(stty_init),
            use_stderr_pipe)

        pcl_expect.Expectable.__init__(self, fd)
        if use_stderr_pipe:
            self.stderr = pcl_expect.Expectable(stderr)

    def send(self, s):
        pcl_expect.debug("sending \"%s\" to fd %d" % (s, self.fileno()))
        os.write(self.fileno(), s)

    def close(self):
        os.close(self.fileno())
        pcl_expect.Expectable.close(self)
        if hasattr(self, "stderr"):
            os.close(self.stderr.fileno())
            self.stderr.close()
        return os.waitpid(self.__child_pid, 0)

def spawn2(cmd, *args, **kwargs):
    proc = Spawn(cmd, use_stderr_pipe=True, *args, **kwargs)
    return proc, proc.stderr