# coding: utf8
# Debian Contributors data mining using subprocesses
#
# Copyright (C) 2014  Enrico Zini <enrico@debian.org>
#
# This program 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 3 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 Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
import codecs
import io
import fcntl
import os
import select
import subprocess

def stream_output(proc):
    """
    Take a subprocess.Popen object and generate its output, as pairs of (tag,
    line) couples. Tag can be O for stdout, E for stderr and R for return
    value.

    Note that the output is not line-split.

    R is always the last bit that gets generated.
    """

    fds = [proc.stdout, proc.stderr]
    tags = ["O", "E"]

    # Set both pipes as non-blocking
    for fd in fds:
        fcntl.fcntl(fd, fcntl.F_SETFL, os.O_NONBLOCK)

    # Multiplex stdout and stderr with different tags
    while len(fds) > 0:
        s = select.select(fds, (), ())
        for fd in s[0]:
            idx = fds.index(fd)
            buf = fd.read()
            if buf:
                decoded = codecs.decode(buf, 'utf-8', 'replace')
                yield tags[idx], decoded
            else:
                fds.pop(idx)
                tags.pop(idx)
    res = proc.wait()
    yield "R", res


class StreamStdoutKeepStderr:
    """
    Stream lines of standard output from a Popen object, keeping all of its
    stderr inside a StringIO
    """
    def __init__(self, proc):
        self.proc = proc
        self.stderr = io.StringIO()

    def __iter__(self):
        last_line = None
        for tag, buf in stream_output(self.proc):
            if tag == "O":
                for l in buf.splitlines(True):
                    if last_line is not None:
                        l = last_line + l
                        last_line = None
                    if l.endswith("\n"):
                        yield l
                    else:
                        last_line = l
            elif tag == "E":
                self.stderr.write(buf)
        if last_line is not None:
            yield last_line

def stream_command_stdout(cmd, **kw):
    try:
        proc = subprocess.Popen(
                cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                stderr=subprocess.PIPE, **kw)
        proc.stdin.close()
        lines = StreamStdoutKeepStderr(proc)
        for line in lines:
            yield line
        result = proc.wait()
    except:
        proc.terminate()
        raise

    if result != 0:
        raise RuntimeError("{} exited with status {}: {}".format(
            cmd[0], result, lines.stderr.getvalue().strip()))

