# Copyright (C) 2008 Gustavo J. A. M. Carneiro  <gjcarneiro@gmail.com>

# 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

import shlex
import subprocess
import sys
import re
import os

env_var_rx = re.compile(r"^([a-zA-Z0-9_]+)=(\S+)$")

def debug(message):
    print >> sys.stderr, message


if sys.platform == 'win32':
    dev_null = open("NUL:", "w")
else:
    dev_null = open("/dev/null", "w")

fcntl = fd = fl = None
try:
    import fcntl
except ImportError:
    pass
else:
    fd = dev_null.fileno()
    fl = fcntl.fcntl(fd, fcntl.F_GETFD)
    fcntl.fcntl(fd, fcntl.F_SETFD, fl | fcntl.FD_CLOEXEC)
del fcntl, fd, fl

def _open_out_file(filename):
    if filename in ['NUL:', '/dev/null']:
        return dev_null
    else:
        return open(filename, 'wb')


class Node(object):
    pass

class Op(Node):
    pass

class Pipe(Op):
    pass

class And(Op):
    pass

class Or(Op):
    pass

class Command(Node):
    class PIPE(object):
        pass # PIPE is a constant
    class STDOUT(object):
        pass # PIPE is a constant

    def __init__(self, name):
        super(Command, self).__init__()
        self.name = name # command name
        self.argv = [name] # command argv
        self.stdin = None
        self.stdout = None
        self.stderr = None
        self.env_vars = None

    def __repr__(self):
        return "Command(%r, argv=%r, stdin=%r, stdout=%r, stderr=%r)" \
            % (self.name, self.argv, self.stdin, self.stdout, self.stderr)

class Chdir(Node):
    def __init__(self):
        super(Chdir, self).__init__()
        self.dir = None

    def __repr__(self):
        return "Chdir(%r)" \
            % (self.dir)

class Pipeline(object):
    def __init__(self):
        self.current_command = None
        self.pipeline = []

    def _commit_command(self):
        assert self.current_command is not None
        self.pipeline.append(self.current_command)
        self.current_command = None

    def get_abbreviated_command(self):
        l = []
        for node in self.pipeline:
            if isinstance(node, Command):
                l.append(node.name)
            if isinstance(node, Chdir):
                l.append('cd %s' % node.dir)
            elif isinstance(node, Pipe):
                l.append('|')
            elif isinstance(node, And):
                l.append('&&')
            elif isinstance(node, And):
                l.append('||')
        return ' '.join(l)

    def parse(self, command):
        self.current_command = None
        self.pipeline = []

        if isinstance(command, list):
            tokens = list(command)
        else:
            tokens = shlex.split(command)
        debug("command: shlex: %r" % (tokens,))

        BEGIN, COMMAND, CHDIR, STDERR, STDOUT, STDIN = range(6)
        state = BEGIN
        self.current_command = None
        env_vars = dict()

        while tokens:
            token = tokens.pop(0)
            if state == BEGIN:
                env_var_match = env_var_rx.match(token)
                if env_var_match is not None:
                    env_vars[env_var_match.group(1)] = env_var_match.group(2)
                else:
                    assert self.current_command is None
                    if token == 'cd':
                        self.current_command = Chdir()
                        assert not env_vars
                        state = CHDIR
                    else:
                        self.current_command = Command(token)
                        if env_vars:
                            self.current_command.env_vars = env_vars
                            env_vars = dict()
                        state = COMMAND
            elif state == COMMAND:
                if token == '>':
                    state = STDOUT
                elif token == '2>':
                    state = STDERR
                elif token == '2>&1':
                    assert self.current_command.stderr is None
                    self.current_command.stderr = Command.STDOUT
                elif token == '<':
                    state = STDIN
                elif token == '|':
                    assert self.current_command.stdout is None
                    self.current_command.stdout = Command.PIPE
                    self._commit_command()
                    self.pipeline.append(Pipe())
                    state = BEGIN
                elif token == '&&':
                    self._commit_command()
                    self.pipeline.append(And())
                    state = BEGIN
                elif token == '||':
                    self._commit_command()
                    self.pipeline.append(Or())
                    state = BEGIN
                else:
                    self.current_command.argv.append(token)
            elif state == CHDIR:
                if token == '&&':
                    self._commit_command()
                    self.pipeline.append(And())
                    state = BEGIN
                else:
                    assert self.current_command.dir is None
                    self.current_command.dir = token
            elif state == STDOUT:
                assert self.current_command.stdout is None
                self.current_command.stdout = token
                state = COMMAND
            elif state == STDERR:
                assert self.current_command.stderr is None
                self.current_command.stderr = token
                state = COMMAND
            elif state == STDIN:
                assert self.current_command.stdin is None
                self.current_command.stdin = token
                state = COMMAND
        self._commit_command()
        return self.pipeline

    def _exec_piped_commands(self, commands):
        retvals = []
        for cmd in commands:
            retvals.append(cmd.wait())
        retval = 0
        for r in retvals:
            if r:
                retval = retvals[-1]
                break
        return retval

    def run(self, verbose=False):
        pipeline = list(self.pipeline)
        files_to_close = []
        piped_commands = []
        piped_commands_display = []
        BEGIN, PIPE = range(2)
        state = BEGIN
        cwd = '.'
        while pipeline:
            node = pipeline.pop(0)

            if isinstance(node, Chdir):
                next_op = pipeline.pop(0)
                assert isinstance(next_op, And)
                cwd = os.path.join(cwd, node.dir)
                if verbose:
                    piped_commands_display.append("cd %s &&" % node.dir)
                continue
            
            assert isinstance(node, (Command, Chdir))
            cmd = node
            if verbose:
                if cmd.env_vars:
                    env_vars_str = ' '.join(['%s=%s' % (key, val) for key, val in cmd.env_vars.iteritems()])
                    piped_commands_display.append("%s %s" % (env_vars_str, ' '.join(cmd.argv)))
                else:
                    piped_commands_display.append(' '.join(cmd.argv))

            if state == PIPE:
                stdin = piped_commands[-1].stdout
            elif cmd.stdin is not None:
                stdin = open(cmd.stdin, "r")
                if verbose:
                    piped_commands_display.append('< %s' % cmd.stdin)
                files_to_close.append(stdin)
            else:
                stdin = None

            if cmd.stdout is None:
                stdout = None
            elif cmd.stdout is Command.PIPE:
                stdout = subprocess.PIPE
            else:
                stdout = _open_out_file(cmd.stdout)
                files_to_close.append(stdout)
                if verbose:
                    piped_commands_display.append('> %s' % cmd.stdout)

            if cmd.stderr is None:
                stderr = None
            elif cmd.stderr is Command.PIPE:
                stderr = subprocess.PIPE
            elif cmd.stderr is Command.STDOUT:
                stderr = subprocess.STDOUT
                if verbose:
                    piped_commands_display.append('2>&1')
            else:
                stderr = _open_out_file(cmd.stderr)
                files_to_close.append(stderr)
                if verbose:
                    piped_commands_display.append('2> %s' % cmd.stderr)

            if cmd.env_vars:
                env = dict(os.environ)
                env.update(cmd.env_vars)
            else:
                env = None

            if cwd == '.':
                proc_cwd = None
            else:
                proc_cwd = cwd

            debug("command: subprocess.Popen(argv=%r, stdin=%r, stdout=%r, stderr=%r, env_vars=%r, cwd=%r)"
                  % (cmd.argv, stdin, stdout, stderr, cmd.env_vars, proc_cwd))
            proc = subprocess.Popen(cmd.argv, stdin=stdin, stdout=stdout, stderr=stderr, env=env, cwd=proc_cwd)
            del stdin, stdout, stderr
            piped_commands.append(proc)

            try:
                next_node = pipeline.pop(0)
            except IndexError:
                try:
                    retval = self._exec_piped_commands(piped_commands)
                    if verbose:
                        print "%s: exit code %i" % (' '.join(piped_commands_display), retval)
                finally:
                    for f in files_to_close:
                        if f is not dev_null:
                            f.close()
                    files_to_close = []
                return retval
            else:

                if isinstance(next_node, Pipe):
                    state = PIPE
                    piped_commands_display.append('|')

                elif isinstance(next_node, Or):
                    try:
                        this_retval = self._exec_piped_commands(piped_commands)
                    finally:
                        for f in files_to_close:
                            if f is not dev_null:
                                f.close()
                        files_to_close = []
                    if this_retval == 0:
                        if verbose:
                            print "%s: exit code %i (|| is short-circuited)" % (' '.join(piped_commands_display), retval)
                        return this_retval
                    if verbose:
                        print "%s: exit code %i (|| proceeds)" % (' '.join(piped_commands_display), retval)
                    state = BEGIN
                    piped_commands = []
                    piped_commands_display = []

                elif isinstance(next_node, And):
                    try:
                        this_retval = self._exec_piped_commands(piped_commands)
                    finally:
                        for f in files_to_close:
                            if f is not dev_null:
                                f.close()
                        files_to_close = []
                    if this_retval != 0:
                        if verbose:
                            print "%s: exit code %i (&& is short-circuited)" % (' '.join(piped_commands_display), retval)
                        return this_retval
                    if verbose:
                        print "%s: exit code %i (&& proceeds)" % (' '.join(piped_commands_display), retval)
                    state = BEGIN
                    piped_commands = []
                    piped_commands_display = []



def _main():
    pipeline = Pipeline()
    pipeline.parse('./foo.py 2>&1 < xxx | cat && ls')
    print pipeline.run()

if __name__ == '__main__':
    _main()