# 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")

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()