from __future__ import with_statement
import six
import subprocess
import time
import atexit
from tempfile import TemporaryFile
from subprocess import Popen, PIPE
from threading import Thread
from plumbum.lib import MinHeap, ascii, bytes
from contextlib import contextmanager


if not hasattr(Popen, "kill"):
    # python 2.5 compatibility
    import os
    import signal
    def _Popen_kill(self):
        os.kill(self.pid, signal.SIGTERM)
    def _Popen_send_signal(self, sig):
        os.kill(self.pid, sig)
    Popen.kill = _Popen_kill
    Popen.terminate = _Popen_kill
    Popen.send_signal = _Popen_send_signal

#===================================================================================================
# Exceptions
#===================================================================================================
class ProcessExecutionError(Exception):
    """Represents the failure of a process. When the exit code of a terminated process does not
    match the expected result, this exception is raised by :func:`run_proc
    <plumbum.commands.run_proc>`. It contains the process' return code, stdout, and stderr, as
    well as the command line used to create the process (``argv``)
    """
    def __init__(self, argv, retcode, stdout, stderr):
        Exception.__init__(self, argv, retcode, stdout, stderr)
        self.argv = argv
        self.retcode = retcode
        if isinstance(stdout, bytes) and not isinstance(stderr, str):
            stdout = ascii(stdout)
        if isinstance(stderr, bytes) and not isinstance(stderr, str):
            stderr = ascii(stderr)
        self.stdout = stdout
        self.stderr = stderr
    def __str__(self):
        stdout = "\n         | ".join(self.stdout.splitlines())
        stderr = "\n         | ".join(self.stderr.splitlines())
        lines = ["Command line: %r" % (self.argv,), "Exit code: %s" % (self.retcode)]
        if stdout:
            lines.append("Stdout:  | %s" % (stdout,))
        if stderr:
            lines.append("Stderr:  | %s" % (stderr,))
        return "\n".join(lines)

class ProcessTimedOut(Exception):
    """Raises by :func:`run_proc <plumbum.commands.run_proc>` when a ``timeout`` has been
    specified and it has elapsed before the process terminated"""
    def __init__(self, msg, argv):
        Exception.__init__(self, msg, argv)
        self.argv = argv

class CommandNotFound(Exception):
    """Raised by :func:`local.which <plumbum.local_machine.LocalMachine.which>` and
    :func:`RemoteMachine.which <plumbum.remote_machine.RemoteMachine.which>` when a
    command was not found in the system's ``PATH``"""
    def __init__(self, program, path):
        Exception.__init__(self, program, path)
        self.program = program
        self.path = path

class RedirectionError(Exception):
    """Raised when an attempt is made to redirect an process' standard handle,
    which was already redirected to/from a file"""


#===================================================================================================
# Utilities
#===================================================================================================
# modified from the stdlib pipes module for windows
_safechars = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@%_-+=:,./'
_funnychars = '"`$\\'
def shquote(text):
    """Quotes the given text with shell escaping (assumes as syntax similar to ``sh``)"""
    if not text:
        return "''"
    text = str(text)
    if not text:
        return "''"
    for c in text:
        if c not in _safechars:
            break
    else:
        return text
    if "'" not in text:
        return "'" + text + "'"
    res = "".join(('\\' + c if c in _funnychars else c) for c in text)
    return '"' + res + '"'

def shquote_list(seq):
    return [shquote(item) for item in seq]

queue = six.moves.queue
_timeout_queue = queue.Queue()
_shutting_down = False

def _timeout_thread():
    waiting = MinHeap()
    while not _shutting_down:
        if waiting:
            ttk, _ = waiting.peek()
            timeout = max(0, ttk - time.time())
        else:
            timeout = None
        try:
            proc, time_to_kill = _timeout_queue.get(timeout = timeout)
            if proc is SystemExit:
                # terminate
                return
            waiting.push((time_to_kill, proc))
        except queue.Empty:
            pass
        now = time.time()
        while waiting:
            ttk, proc = waiting.peek()
            if ttk > now:
                break
            waiting.pop()
            try:
                if proc.poll() is None:
                    proc.kill()
                    proc._timed_out = True
            except EnvironmentError:
                pass

thd1 = Thread(target = _timeout_thread, name = "PlumbumTimeoutThread")
thd1.setDaemon(True)
thd1.start()

def _shutdown_bg_threads():
    global _shutting_down
    _shutting_down = True
    _timeout_queue.put((SystemExit, 0))

atexit.register(_shutdown_bg_threads)

def run_proc(proc, retcode, timeout = None):
    """Waits for the given process to terminate, with the expected exit code

    :param proc: a running Popen-like object

    :param retcode: the expected return (exit) code of the process. It defaults to 0 (the
                    convention for success). If ``None``, the return code is ignored.
                    It may also be a tuple (or any object that supports ``__contains__``)
                    of expected return codes.

    :param timeout: the number of seconds (a ``float``) to allow the process to run, before
                    forcefully terminating it. If ``None``, not timeout is imposed; otherwise
                    the process is expected to terminate within that timeout value, or it will
                    be killed and :class:`ProcessTimedOut <plumbum.cli.ProcessTimedOut>`
                    will be raised

    :returns: A tuple of (return code, stdout, stderr)
    """
    if timeout is not None:
        _timeout_queue.put((proc, time.time() + timeout))
    stdout, stderr = proc.communicate()
    proc._end_time = time.time()
    if not stdout:
        stdout = six.b("")
    if not stderr:
        stderr = six.b("")
    if getattr(proc, "encoding", None):
        stdout = stdout.decode(proc.encoding, "ignore")
        stderr = stderr.decode(proc.encoding, "ignore")

    if getattr(proc, "_timed_out", False):
        raise ProcessTimedOut("Process did not terminate within %s seconds" % (timeout,),
            getattr(proc, "argv", None))

    if retcode is not None:
        if hasattr(retcode, "__contains__"):
            if proc.returncode not in retcode:
                raise ProcessExecutionError(getattr(proc, "argv", None), proc.returncode,
                    stdout, stderr)
        elif proc.returncode != retcode:
            raise ProcessExecutionError(getattr(proc, "argv", None), proc.returncode,
                stdout, stderr)
    return proc.returncode, stdout, stderr

#===================================================================================================
# Commands
#===================================================================================================
class BaseCommand(object):
    """Base of all command objects"""

    __slots__ = ["cwd", "env", "encoding"]

    def __str__(self):
        return " ".join(self.formulate())

    def __or__(self, other):
        """Creates a pipe with the other command"""
        return Pipeline(self, other)

    def __gt__(self, file):
        """Redirects the process' stdout to the given file"""
        return StdoutRedirection(self, file)

    def __rshift__(self, file):
        """Redirects the process' stdout to the given file (appending)"""
        return AppendingStdoutRedirection(self, file)

    def __ge__(self, file):
        """Redirects the process' stderr to the given file"""
        return StderrRedirection(self, file)

    def __lt__(self, file):
        """Redirects the given file into the process' stdin"""
        return StdinRedirection(self, file)
    def __lshift__(self, data):
        """Redirects the given data into the process' stdin"""
        return StdinDataRedirection(self, data)

    def __getitem__(self, args):
        """Creates a bound-command with the given arguments"""
        if not isinstance(args, (tuple, list)):
            args = (args,)
        if not args:
            return self
        if isinstance(self, BoundCommand):
            return BoundCommand(self.cmd, self.args + tuple(args))
        else:
            return BoundCommand(self, args)

    def __call__(self, *args, **kwargs):
        """A shortcut for `run(args)`, returning only the process' stdout"""
        return self.run(args, **kwargs)[1]

    def _get_encoding(self):
        raise NotImplementedError()

    def formulate(self, level = 0, args = ()):
        """Formulates the command into a command-line, i.e., a list of shell-quoted strings
        that can be executed by ``Popen`` or shells.

        :param level: The nesting level of the formulation; it dictates how much shell-quoting
                      (if any) should be performed

        :param args: The arguments passed to this command (a tuple)

        :returns: A list of strings
        """
        raise NotImplementedError()

    def popen(self, args = (), **kwargs):
        """Spawns the given command, returning a ``Popen``-like object.

        .. note::

           When processes run in the **background** (either via ``popen`` or
           :class:`& BG <plumbum.commands.BG>`), their stdout/stderr pipes might fill up,
           causing them to hang. If you know a process produces output, be sure to consume it
           every once in a while, using a monitoring thread/reactor in the background.
           For more info, see `#48 <https://github.com/tomerfiliba/plumbum/issues/48>`_

        :param args: Any arguments to be passed to the process (a tuple)

        :param kwargs: Any keyword-arguments to be passed to the ``Popen`` constructor

        :returns: A ``Popen``-like object
        """
        raise NotImplementedError()

    @contextmanager
    def bgrun(self, args = (), **kwargs):
        """Runs the given command as a context manager, allowing you to create a
        `pipeline <http://en.wikipedia.org/wiki/Pipeline_(computing)>`_ (not in the UNIX sense)
        of programs, parallelizing their work. In other words, instead of running programs
        one after the other, you can start all of them at the same time and wait for them to
        finish. For a more thorough review, see
        `Lightweight Asynchronism <http://tomerfiliba.com/blog/Toying-with-Context-Managers/>`_.

        Example::

            from plumbum.cmd import mkfs

            with mkfs["-t", "ext3", "/dev/sda1"] as p1:
                with mkfs["-t", "ext3", "/dev/sdb1"] as p2:
                    pass

        .. note::

           When processes run in the **background** (either via ``popen`` or
           :class:`& BG <plumbum.commands.BG>`), their stdout/stderr pipes might fill up,
           causing them to hang. If you know a process produces output, be sure to consume it
           every once in a while, using a monitoring thread/reactor in the background.
           For more info, see `#48 <https://github.com/tomerfiliba/plumbum/issues/48>`_

        For the arguments, see :func:`run <BaseCommand.run>`.

        :returns: A Popen object, augmented with a ``.run()`` method, which returns a tuple of
                  (return code, stdout, stderr)
        """
        retcode = kwargs.pop("retcode", 0)
        timeout = kwargs.pop("timeout", None)
        p = self.popen(args, **kwargs)
        was_run = [False]
        def runner():
            if was_run[0]:
                return  # already done
            was_run[0] = True
            try:
                return run_proc(p, retcode, timeout)
            finally:
                for f in [p.stdin, p.stdout, p.stderr]:
                    try:
                        f.close()
                    except Exception:
                        pass
        p.run = runner
        yield p
        runner()

    def run(self, args = (), **kwargs):
        """Runs the given command (equivalent to popen() followed by
        :func:`run_proc <plumbum.commands.run_proc>`). If the exit code of the process does
        not match the expected one, :class:`ProcessExecutionError
        <plumbum.commands.ProcessExecutionError>` is raised.

        :param args: Any arguments to be passed to the process (a tuple)

        :param retcode: The expected return code of this process (defaults to 0).
                        In order to disable exit-code validation, pass ``None``. It may also
                        be a tuple (or any iterable) of expected exit codes.

                        .. note:: this argument must be passed as a keyword argument.

        :param timeout: The maximal amount of time (in seconds) to allow the process to run.
                        ``None`` means no timeout is imposed; otherwise, if the process hasn't
                        terminated after that many seconds, the process will be forcefully
                        terminated an exception will be raised

                        .. note:: this argument must be passed as a keyword argument.

        :param kwargs: Any keyword-arguments to be passed to the ``Popen`` constructor

        :returns: A tuple of (return code, stdout, stderr)
        """
        with self.bgrun(args, **kwargs) as p:
            return p.run()


class BoundCommand(BaseCommand):
    __slots__ = ["cmd", "args"]
    def __init__(self, cmd, args):
        self.cmd = cmd
        self.args = args
    def __repr__(self):
        return "BoundCommand(%r, %r)" % (self.cmd, self.args)
    def _get_encoding(self):
        return self.cmd._get_encoding()
    def formulate(self, level = 0, args = ()):
        return self.cmd.formulate(level + 1, self.args + tuple(args))
    def popen(self, args = (), **kwargs):
        if isinstance(args, six.string_types):
            args = (args,)
        return self.cmd.popen(self.args + tuple(args), **kwargs)

class Pipeline(BaseCommand):
    __slots__ = ["srccmd", "dstcmd"]
    def __init__(self, srccmd, dstcmd):
        self.srccmd = srccmd
        self.dstcmd = dstcmd
    def __repr__(self):
        return "Pipeline(%r, %r)" % (self.srccmd, self.dstcmd)
    def _get_encoding(self):
        return self.srccmd._get_encoding() or self.dstcmd._get_encoding()
    def formulate(self, level = 0, args = ()):
        return self.srccmd.formulate(level + 1) + ["|"] + self.dstcmd.formulate(level + 1, args)

    def popen(self, args = (), **kwargs):
        src_kwargs = kwargs.copy()
        src_kwargs["stdout"] = PIPE
        src_kwargs["stderr"] = PIPE

        srcproc = self.srccmd.popen(args, **src_kwargs)
        kwargs["stdin"] = srcproc.stdout
        dstproc = self.dstcmd.popen(**kwargs)
        # allow p1 to receive a SIGPIPE if p2 exits
        srcproc.stdout.close()
        srcproc.stderr.close()
        if srcproc.stdin:
            srcproc.stdin.close()
        dstproc.srcproc = srcproc
        return dstproc

class BaseRedirection(BaseCommand):
    __slots__ = ["cmd", "file"]
    SYM = None
    KWARG = None
    MODE = None

    def __init__(self, cmd, file):
        self.cmd = cmd
        self.file = file
    def _get_encoding(self):
        return self.cmd._get_encoding()
    def __repr__(self):
        return "%s(%r, %r)" % (self.__class__.__name__, self.cmd, self.file)
    def formulate(self, level = 0, args = ()):
        return self.cmd.formulate(level + 1, args) + [self.SYM, shquote(getattr(self.file, "name", self.file))]
    def popen(self, args = (), **kwargs):
        from plumbum.local_machine import LocalPath
        from plumbum.remote_machine import RemotePath

        if self.KWARG in kwargs and kwargs[self.KWARG] not in (PIPE, None):
            raise RedirectionError("%s is already redirected" % (self.KWARG,))
        if isinstance(self.file, six.string_types + (LocalPath,)):
            f = kwargs[self.KWARG] = open(str(self.file), self.MODE)
        elif isinstance(self.file, RemotePath):
            raise TypeError("Cannot redirect to/from remote paths")
        else:
            kwargs[self.KWARG] = self.file
            f = None
        try:
            return self.cmd.popen(args, **kwargs)
        finally:
            if f:
                f.close()

class StdinRedirection(BaseRedirection):
    __slots__ = []
    SYM = "<"
    KWARG = "stdin"
    MODE = "r"

class StdoutRedirection(BaseRedirection):
    __slots__ = []
    SYM = ">"
    KWARG = "stdout"
    MODE = "w"

class AppendingStdoutRedirection(BaseRedirection):
    __slots__ = []
    SYM = ">>"
    KWARG = "stdout"
    MODE = "a"

class StderrRedirection(BaseRedirection):
    __slots__ = []
    SYM = "2>"
    KWARG = "stderr"
    MODE = "w"

class ERROUT(int):
    def __repr__(self):
        return "ERROUT"
    def __str__(self):
        return "&1"
ERROUT = ERROUT(subprocess.STDOUT)

class StdinDataRedirection(BaseCommand):
    __slots__ = ["cmd", "data"]
    CHUNK_SIZE = 16000

    def __init__(self, cmd, data):
        self.cmd = cmd
        self.data = data
    def _get_encoding(self):
        return self.cmd._get_encoding()

    def formulate(self, level = 0, args = ()):
        return ["echo %s" % (shquote(self.data),), "|", self.cmd.formulate(level + 1, args)]
    def popen(self, args = (), **kwargs):
        if "stdin" in kwargs and kwargs["stdin"] != PIPE:
            raise RedirectionError("stdin is already redirected")
        data = self.data
        if not isinstance(data, bytes) and self._get_encoding() is not None:
            data = data.encode(self._get_encoding())
        f = TemporaryFile()
        while data:
            chunk = data[:self.CHUNK_SIZE]
            f.write(chunk)
            data = data[self.CHUNK_SIZE:]
        f.seek(0)
        # try:
        return self.cmd.popen(args, stdin = f, **kwargs)
        # finally:
        #    f.close()

class ConcreteCommand(BaseCommand):
    QUOTE_LEVEL = None
    __slots__ = ["executable", "encoding"]
    def __init__(self, executable, encoding):
        self.executable = executable
        self.encoding = encoding
        self.cwd = None
        self.env = None
    def __str__(self):
        return str(self.executable)
    def _get_encoding(self):
        return self.encoding

    def formulate(self, level = 0, args = ()):
        argv = [str(self.executable)]
        for a in args:
            if not a and a != "":
                continue
            if isinstance(a, BaseCommand):
                if level >= self.QUOTE_LEVEL:
                    argv.extend(shquote_list(a.formulate(level + 1)))
                else:
                    argv.extend(a.formulate(level + 1))
            else:
                if level >= self.QUOTE_LEVEL:
                    argv.append(shquote(a))
                else:
                    argv.append(str(a))
        # if self.encoding:
        #    argv = [a.encode(self.encoding) for a in argv if isinstance(a, six.string_types)]
        return argv

#===================================================================================================
# execution modifiers (background, foreground)
#===================================================================================================
class ExecutionModifier(object):
    __slots__ = ["retcode"]
    def __init__(self, retcode = 0):
        self.retcode = retcode
    def __repr__(self):
        return "%s(%r)" % (self.__class__.__name__, self.retcode)
    @classmethod
    def __call__(cls, retcode):
        return cls(retcode)

class Future(object):
    """Represents a "future result" of a running process. It basically wraps a ``Popen``
    object and the expected exit code, and provides poll(), wait(), returncode, stdout,
    and stderr.
    """
    def __init__(self, proc, expected_retcode, timeout = None):
        self.proc = proc
        self._expected_retcode = expected_retcode
        self._timeout = timeout
        self._returncode = None
        self._stdout = None
        self._stderr = None
    def __repr__(self):
        return "<Future %r (%s)>" % (self.proc.argv, self._returncode if self.ready() else "running",)
    def poll(self):
        """Polls the underlying process for termination; returns ``None`` if still running,
        or the process' returncode if terminated"""
        if self.proc.poll() is not None:
            self.wait()
        return self._returncode is not None
    ready = poll
    def wait(self):
        """Waits for the process to terminate; will raise a
        :class:`plumbum.commands.ProcessExecutionError` in case of failure"""
        if self._returncode is not None:
            return
        self._returncode, self._stdout, self._stderr = run_proc(self.proc,
            self._expected_retcode, self._timeout)
    @property
    def stdout(self):
        """The process' stdout; accessing this property will wait for the process to finish"""
        self.wait()
        return self._stdout
    @property
    def stderr(self):
        """The process' stderr; accessing this property will wait for the process to finish"""
        self.wait()
        return self._stderr
    @property
    def returncode(self):
        """The process' returncode; accessing this property will wait for the process to finish"""
        self.wait()
        return self._returncode

class BG(ExecutionModifier):
    """
    An execution modifier that runs the given command in the background, returning a
    :class:`Future <plumbum.commands.Future>` object. In order to mimic shell syntax, it applies
    when you right-and it with a command. If you wish to expect a different return code
    (other than the normal success indicate by 0), use ``BG(retcode)``. Example::

        future = sleep[5] & BG       # a future expecting an exit code of 0
        future = sleep[5] & BG(7)    # a future expecting an exit code of 7

    .. note::

       When processes run in the **background** (either via ``popen`` or
       :class:`& BG <plumbum.commands.BG>`), their stdout/stderr pipes might fill up,
       causing them to hang. If you know a process produces output, be sure to consume it
       every once in a while, using a monitoring thread/reactor in the background.
       For more info, see `#48 <https://github.com/tomerfiliba/plumbum/issues/48>`_
    """
    __slots__ = []
    def __rand__(self, cmd):
        return Future(cmd.popen(), self.retcode)

BG = BG()
"""
An execution modifier that runs the given command in the background, returning a
:class:`Future <plumbum.commands.Future>` object. In order to mimic shell syntax, it applies
when you right-and it with a command. If you wish to expect a different return code
(other than the normal success indicate by 0), use ``BG(retcode)``. Example::

    future = sleep[5] & BG       # a future expecting an exit code of 0
    future = sleep[5] & BG(7)    # a future expecting an exit code of 7

.. note::

   When processes run in the **background** (either via ``popen`` or
   :class:`& BG <plumbum.commands.BG>`), their stdout/stderr pipes might fill up,
   causing them to hang. If you know a process produces output, be sure to consume it
   every once in a while, using a monitoring thread/reactor in the background.
   For more info, see `#48 <https://github.com/tomerfiliba/plumbum/issues/48>`_
"""

class FG(ExecutionModifier):
    """
    An execution modifier that runs the given command in the foreground, passing it the
    current process' stdin, stdout and stderr. Useful for interactive programs that require
    a TTY. There is no return value.

    In order to mimic shell syntax, it applies when you right-and it with a command.
    If you wish to expect a different return code (other than the normal success indicate by 0),
    use ``BG(retcode)``. Example::

        vim & FG       # run vim in the foreground, expecting an exit code of 0
        vim & FG(7)    # run vim in the foreground, expecting an exit code of 7
    """
    __slots__ = []
    def __rand__(self, cmd):
        cmd(retcode = self.retcode, stdin = None, stdout = None, stderr = None)

FG = FG()
"""
An execution modifier that runs the given command in the foreground, passing it the
current process' stdin, stdout and stderr. Useful for interactive programs that require
a TTY. There is no return value.

In order to mimic shell syntax, it applies when you right-and it with a command.
If you wish to expect a different return code (other than the normal success indicate by 0),
use ``BG(retcode)``. Example::

    vim & FG       # run vim in the foreground, expecting an exit code of 0
    vim & FG(7)    # run vim in the foreground, expecting an exit code of 7
"""

class Tee(object):
    def __init__(self, *streams):
        self.streams = streams












