#
#

# Copyright (C) 2009 Google Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program 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
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.


"""Module with helper classes and functions for daemons.

"""


import gobject
import logging
import os

_PROCESS_EXIT_IO_TIMEOUT = 1


class SignalRegistration:
  def __init__(self, emitter, handle):
    self.__emitter = emitter
    self.__handle = handle

  def __del__(self):
    if self.__emitter is not None:
      self.Disconnect()

  def Disconnect(self):
    assert self.__emitter is not None
    assert self.__handle is not None

    self.__emitter.disconnect(self.__handle)

    self.__emitter = None
    self.__handle = None


class IOChannel(gobject.GObject, object):
  """Wrapper for gobject.IOChannel.

  Emits signals on I/O events.

  """
  AFTER_READ_SIGNAL = "after-read"
  WRITE_COMPLETE_SIGNAL = "after-write"
  CLOSED_SIGNAL = "closed"

  __BLOCKSIZE = 1024

  __gsignals__ = {
    AFTER_READ_SIGNAL:
      (gobject.SIGNAL_RUN_LAST, gobject.TYPE_NONE,
       (gobject.TYPE_PYOBJECT, )),
    WRITE_COMPLETE_SIGNAL:
      (gobject.SIGNAL_RUN_LAST, gobject.TYPE_NONE,
       ()),
    CLOSED_SIGNAL:
      (gobject.SIGNAL_RUN_LAST, gobject.TYPE_NONE,
       ()),
    }

  def __init__(self):
    """Initializes this class.

    """
    object.__init__(self)
    gobject.GObject.__init__(self)
    #logging.debug("Creating IO Channel: %r", self)
    self.__channel = None
    self.__handle = None
    self.__writebuf = ""
    self.__writepos = 0

  def __GetClosed(self):
    """Returns whether this channel has been closed.

    """
    return self.__channel is None

  closed = property(fget=__GetClosed)

  @staticmethod
  def __CreateChannel(fd):
    """Creates the inner gobject.IOChannel instance.

    Note: has side-effects on file descriptor by setting O_NONBLOCK.

    @type fd: int
    @param fd: File descriptor

    """
    channel = gobject.IOChannel(fd)
    channel.set_flags(gobject.IO_FLAG_NONBLOCK)
    #channel.set_encoding(None)
    #channel.set_buffered(0)
    channel.set_encoding("UTF-8")
    return channel

  def Attach(self, fd):
    """Attaches this channel to a file descriptor.

    @type fd: int
    @param fd: File descriptor

    """
    #logging.debug("Attached fd %r to %r", fd, self)
    assert self.__channel is None
    self.__channel = self.__CreateChannel(fd)
    self.__fd = fd
    self.__Update(False)

  def Detach(self):
    """Closes this channel.

    """
    if self.__channel:
      self.__Close()
    else:
      self.__Update(True)

  def Write(self, data):
    """Asynchronous write.

    @type data: str
    @param data: Data to be written

    """
    #logging.debug("%r attempting async write or %r",self,data)
    self.__writebuf = self.__writebuf[self.__writepos:] + data
    self.__writepos = 0
    self.__Update(False)

  def __Update(self, detach):
    """Updates this channel's registration with the mainloop.

    @type detach: bool
    @param detach: Whether we're detaching from the file descriptor

    """
    handle = self.__handle

    # TODO: Don't remove and re-add if conditions didn't change
    if handle is not None:
      gobject.source_remove(handle)
      handle = None

    if self.__channel and not detach:
      condition = self.__CalcCondition()
      if condition:
        handle = self.__channel.add_watch(condition, self.__HandleIO)

    self.__handle = handle

  def __CalcCondition(self):
    """Returns necessary flags for mainloop.

    """
    cond = gobject.IO_IN | gobject.IO_HUP | gobject.IO_ERR | gobject.IO_NVAL
    #logging.debug("calculating condition of %r", self)
    if self.__writebuf and self.__writepos < len(self.__writebuf):
      #logging.debug("decided that %r needs IO_OUT", self)
      cond |= gobject.IO_OUT
    #logging.debug("condition of %r is:%r", self, cond)
    return cond

  def __Read(self, channel):
    """Reads data from channel and emits events.

    """
    #logging.debug("In __Read of %r", self)
    assert channel == self.__channel
    #logging.debug("%r is reading", self)
    data = self.__channel.read(self.__BLOCKSIZE)
    #if(data[-1] != "\x00"):
      #data = data +"\x00"
    if data:
      #logging.debug("%r read %r bytes, %r", self, len(data), data)
      self.__EmitAfterRead(data)
      return True
    #logging.debug("%r read no data", self)
    # TODO: Correct when still writing?
    self.__Close()
    return False

  def __Write(self, channel):
    """Writes data from the buffer to the channel.

    """
    #logging.debug("In __Write of %r", self)
    assert channel == self.__channel
    #logging.debug("Passed assert for __Write of %r", self)
    endpos = self.__writepos + self.__BLOCKSIZE

    # Use slicing to avoid unnecessary memory reallocations
    try:
      data = str(self.__writebuf[self.__writepos:endpos], encoding="UTF-8")
    except:
      #logging.debug("excepted in %r for type %r",self, type(self.__writebuf))
      pass
    try:
      data = str(bytes(self.__writebuf[self.__writepos:endpos], encoding="UTF-8"), encoding="UTF-8")
    except:
      #logging.debug("excepted in %r for type %r",self, type(self.__writebuf))
      pass
      
    #logging.debug("%r is writing %r",self,data)
    n = os.write(self.__fd, data.encode("UTF-8"))
    #logging.debug("passed the write in %r", self)
    if n == 0:
      self.__Close()
      return False
    #logging.debug("%r wrote %r bytes",self, n)
    self.__writepos += n

    assert self.__writepos <= len(self.__writebuf)

    if self.__writepos == len(self.__writebuf):
      self.__Update(False)
      self.__EmitWriteComplete()

      # TODO: Cleaner solution
      if self.__writepos == len(self.__writebuf):
        self.__Close()
        return False

    return True

  def __HandleIO(self, channel, cond):
    """Triages I/O events.

    """
    assert channel == self.__channel

    if cond & (gobject.IO_IN | gobject.IO_OUT):
      #logging.debug("Condition in __HandleIO passed for %r", self)
      return (((cond & gobject.IO_IN) and self.__Read(channel) or
              ((cond & gobject.IO_OUT) and self.__Write(channel))))

    if cond & (gobject.IO_HUP | gobject.IO_ERR | gobject.IO_NVAL):
      self.__Close()

    return False

  def __Close(self):
    """Closes the channel.

    """
    #logging.debug("Closed %r", self)
    self.__channel.close(flush=True)
    self.__channel = None
    self.__Update(True)
    self.__EmitClosed()

  def __EmitAfterRead(self, data):
    #logging.debug("In __EmitAfterRead for %r", self)
    self.emit(self.AFTER_READ_SIGNAL, data)

  def __EmitWriteComplete(self):
    self.emit(self.WRITE_COMPLETE_SIGNAL)

  def __EmitClosed(self):
    assert self.closed

    self.emit(self.CLOSED_SIGNAL)


class ChopReader(gobject.GObject):
  """Reads slices separated by separator from L{IOChannel}.

  For each slice, a signal is emitted.

  """
  SLICE_COMPLETE_SIGNAL = "slice-complete"

  __gsignals__ = {
    SLICE_COMPLETE_SIGNAL:
      (gobject.SIGNAL_RUN_LAST, gobject.TYPE_NONE,
       (gobject.TYPE_PYOBJECT, )),
    }

  def __init__(self, sep):
    """Initializes this class.

    """
    #logging.debug("Creating chop reader: %r", self)
    gobject.GObject.__init__(self)
    self.__sep = sep
    self.__channel = None
    self.__after_read_reg = None
    self.__closed_reg = None
    self.__buf = ""

  def Attach(self, channel):
    """Attach to I/O channel.

    @type channel: L{IOChannel}
    @param channel: I/O channel

    """
    assert self.__channel is None
    assert self.__after_read_reg is None
    assert self.__closed_reg is None

    self.__channel = channel

    self.__after_read_reg = \
      SignalRegistration(channel,
                         channel.connect(IOChannel.AFTER_READ_SIGNAL,
                                         self.__ReceivedData))
    self.__closed_reg = \
      SignalRegistration(channel,
                         channel.connect(IOChannel.CLOSED_SIGNAL,
                                         self.__Closed))

  def Detach(self):
    """Detaches from I/O channel.

    """
    self.__channel = None

    if self.__after_read_reg:
      self.__after_read_reg.Disconnect()
      self.__after_read_reg = None

    if self.__closed_reg:
      self.__closed_reg.Disconnect()
      self.__closed_reg = None

  def __del__(self):
    self.Detach()

  def __ParseBuffer(self):
    """Parses internal buffer until no more separators are found.

    """
    #logging.debug("In __ParseBuffer of %r for %r", self, self.__buf)
    pos = 0
    try:
      while True:
        #logging.debug("looking for %r in %r", self.__sep, self.__buf)
        idx = self.__buf.find(self.__sep, pos)
        if idx < 0:
          break

        slice_ = self.__buf[pos:idx]
        pos = idx + len(self.__sep)

        self.__EmitSliceComplete(slice_)
    finally:
      self.__buf = self.__buf[pos:]
    #logging.debug("Escaped __Parsebuffer of %r", self)

  def __ReceivedData(self, channel, data):
    """Adds received data to buffer.

    """
    #logging.debug("In __ReceivedData of %r", self)
    assert channel == self.__channel
    #logging.debug("Passed assert")
    if data:
      self.__buf += str(data, encoding="UTF-8")
    #logging.debug("Going to __ParseBuffer")
    self.__ParseBuffer()
    #logging.debug("Leaving __ReceivedData of %r", self)

  def __Closed(self, channel):
    """Handles leftovers in buffer.

    """
    assert channel == self.__channel
    #logging.debug("In closed of %r", self)
    self.__ParseBuffer()

    # Handle rest
    if self.__buf:
      if self.__buf.endswith(self.__sep):
        self.__buf = self.__buf[:-len(self.__sep)]
      self.__EmitSliceComplete(self.__buf)
      self.__buf = ""

  def __EmitSliceComplete(self, slice_):
    self.emit(self.SLICE_COMPLETE_SIGNAL, slice_)


class Program(gobject.GObject, object):
  """Wrapper for gobject.spawn_async.

  Emits signals on events.

  """
  EXITED_SIGNAL = "exited"

  __gsignals__ = {
    EXITED_SIGNAL:
      (gobject.SIGNAL_RUN_LAST, gobject.TYPE_NONE,
       (gobject.TYPE_PYOBJECT, gobject.TYPE_PYOBJECT)),
    }

  def __init__(self, args, env=None, cwd=None, executable=None,
               umask=None, stdin_data=None):
    """Initializes this class.

    @type args: list
    @param args: Program arguments
    @type env: dict
    @param env: Environment variables for program
    @type cwd: str
    @param cwd: Working directory for program
    @type executable: str
    @param executable: If set, the executable to run
    @type stdin_data: str
    @param stdin_data: Data to be written to program's stdin

    """
    object.__init__(self)
    gobject.GObject.__init__(self)

    self.__args = args
    self.__env = env
    self.__cwd = cwd
    self.__executable = executable
    self.__umask = umask
    self.__stdin_data = stdin_data

    self.__progname = args[0]

    self.__pid = None
    self.__exitcode = None
    self.__child_watch_handle = None

    self.stdin = IOChannel()
    self.__stdin_closed_reg = \
      SignalRegistration(self.stdin,
                         self.stdin.connect(IOChannel.CLOSED_SIGNAL,
                                            self.__HandlePipeClosed))

    self.stdout = IOChannel()
    self.__stdout_closed_reg = \
      SignalRegistration(self.stdout,
                         self.stdout.connect(IOChannel.CLOSED_SIGNAL,
                                             self.__HandlePipeClosed))

    self.stderr = IOChannel()
    self.__stderr_closed_reg = \
      SignalRegistration(self.stderr,
                         self.stderr.connect(IOChannel.CLOSED_SIGNAL,
                                             self.__HandlePipeClosed))

    self.stdout_line = ChopReader(os.linesep)
    self.__stdout_line_complete_reg = \
      SignalRegistration(self.stdout_line,
                         self.stdout_line.connect(ChopReader.SLICE_COMPLETE_SIGNAL,
                                                  self.__LogOutput, "stdout"))
    self.stdout_line.Attach(self.stdout)

    self.stderr_line = ChopReader(os.linesep)
    self.__stderr_line_complete_reg = \
      SignalRegistration(self.stderr_line,
                         self.stderr_line.connect(ChopReader.SLICE_COMPLETE_SIGNAL,
                                                  self.__LogOutput, "stderr"))
    self.stderr_line.Attach(self.stderr)
    #logging.debug("Program: %r, data: %r", self.__progname, stdin_data)
    if stdin_data:
      self.stdin.Write(stdin_data)

  def __GetPid(self):
    """Returns the PID of the started program.

    """
    return self.__pid

  pid = property(fget=__GetPid)

  def __LogOutput(self, _, line, pipename):
    logging.debug("%s[%d] %s: %s", self.__progname, self.pid, pipename, line)

  @staticmethod
  def __FormatEnvironment(env):
    """Formats environment as requested by gobject.spawn_async.

    """
    # gobject.spawn_async needs this list to be a plain string, not unicode
    return [str("%s=%s" % (key, value)) for (key, value) in list(env.items())]

  @staticmethod
  def __FormatArgs(args):
    """Formats arguments as requested by gobject.spawn_async.

    """
    # gobject.spawn_async needs this list to be a plain string, not unicode
    return list(map(str, args))

  def __ChildSetup(self):
    """Called in child process just before the actual program is executed.

    """
    if self.__umask is not None:
      os.umask(self.__umask)

  def Start(self):
    """Start program.

    """
    logging.info("Starting program, executable=%r, args=%r",
                 self.__executable, self.__args)

    # TODO: Make gobject.SPAWN_SEARCH_PATH controllable by caller?
    flags = gobject.SPAWN_DO_NOT_REAP_CHILD | gobject.SPAWN_SEARCH_PATH

    if self.__executable:
      args = self.__FormatArgs([self.__executable] + self.__args)
      flags |= gobject.SPAWN_FILE_AND_ARGV_ZERO
    else:
      args = self.__FormatArgs(self.__args)

    # gobject.spawn_async doesn't take None for default values, hence we can
    # only fill these parameters if they need to be set.
    kwargs = {}

    if self.__env is not None:
      kwargs["envp"] = self.__FormatEnvironment(self.__env)

    if self.__cwd is not None:
      kwargs["working_directory"] = self.__cwd

    (pid, stdin_fd, stdout_fd, stderr_fd) = \
      gobject.spawn_async(args, flags=flags, child_setup=self.__ChildSetup,
                          standard_input=True,
                          standard_output=True,
                          standard_error=True,
                          **kwargs)

    logging.info("Child %s[%d] started", self.__progname, pid)
    self.__pid = pid

    self.stdin.Attach(stdin_fd)
    self.__stdin_fd = stdin_fd
    self.stdout.Attach(stdout_fd)
    self.stderr.Attach(stderr_fd)

    self.__child_watch_handle = \
      gobject.child_watch_add(self.__pid, self.__HandleExit)

    return self.pid

  def __HandleExit(self, pid, exitcode):
    """Called when program exits.

    @type pid: int
    @param pid: Process ID
    @type exitcode: int
    @param exitcode: Exit status as returned by waitpid(2)

    """
    assert pid == self.__pid
    assert self.__exitcode is None

    self.__exitcode = exitcode
    self.__CheckExit()

    # TODO: Should child watch handle be removed from mainloop?

  def __HandlePipeClosed(self, _):
    """Called when I/O pipe to process is closed.

    """
    #logging.debug("%r, closed pipe", self)
    pass

  def __CheckExit(self):
    """Checks whether program has exitted and all I/O has been handled.

    """
    logging.debug("__CheckExit %s[%d] stdin:%r stdout:%r stderr:%r exit:%r",
        self.__progname, self.pid, self.stdin.closed,
        self.stdout.closed, self.stderr.closed, self.__exitcode)

    if (self.stdin.closed and
        self.stdout.closed and
        self.stderr.closed and
        self.__exitcode is not None):
      # TODO: Keep part of stderr output in case of errors (deque)
      exitstatus = None
      signum = None

      if os.WIFSIGNALED(self.__exitcode):
        signum = os.WTERMSIG(self.__exitcode)

      elif os.WIFEXITED(self.__exitcode):
        exitstatus = os.WEXITSTATUS(self.__exitcode)

      else:
        raise RuntimeError("Invalid child status")

      if exitstatus == 0 and signum is None:
        logging.debug("%s[%d] exited cleanly", self.__progname, self.pid)
      else:
        logging.error("%s[%d] failed (status=%s, signal=%s)",
                      self.__progname, self.pid, exitstatus, signum)

      self.__EmitExited(exitstatus, signum)

    else:
      # Add timeout to close stdin/out/err
      gobject.timeout_add(_PROCESS_EXIT_IO_TIMEOUT * 1000,
                          self.__CloseAndExit)


  def __CloseAndExit(self):
    logging.debug("__CloseEndExit %s[%d] stdin:%r stdout:%r stderr:%r",
        self.__progname, self.pid, self.stdin.closed,
        self.stdout.closed, self.stderr.closed)
    self.stdin.Detach()
    self.stdout.Detach()
    self.stderr.Detach()
    self.__CheckExit()

  def __EmitExited(self, exitcode, signum):
    self.emit(self.EXITED_SIGNAL, exitcode, signum)
