from __future__ import unicode_literals

from socket import timeout as timeout_error
from subprocess import PIPE, Popen

import click
import gevent

from pyinfra import logger
from pyinfra.api.util import make_command, read_buffer


def connect(state, host, **kwargs):
    return True


def run_shell_command(
    state, host, command,
    sudo=False, sudo_user=None, su_user=None, preserve_sudo_env=False,
    get_pty=False, env=None, timeout=None, print_output=False,
):
    command = make_command(
        command,
        env=env,
        sudo=sudo,
        sudo_user=sudo_user,
        su_user=su_user,
        preserve_sudo_env=preserve_sudo_env,
    )

    logger.debug('--> Running command on localhost: {0}'.format(command))

    if print_output:
        print('{0}>>> {1}'.format(host.print_prefix, command))

    process = Popen(command, shell=True, stdout=PIPE, stderr=PIPE)

    # Iterate through outputs to get an exit status and generate desired list
    # output, done in two greenlets so stdout isn't printed before stderr. Not
    # attached to state.pool to avoid blocking it with 2x n-hosts greenlets.
    stdout_reader = gevent.spawn(
        read_buffer, process.stdout,
        print_output=print_output,
        print_func=lambda line: '{0}{1}'.format(host.print_prefix, line),
    )
    stderr_reader = gevent.spawn(
        read_buffer, process.stderr,
        print_output=print_output,
        print_func=lambda line: '{0}{1}'.format(
            host.print_prefix, click.style(line, 'red'),
        ),
    )

    # Wait on output, with our timeout (or None)
    greenlets = gevent.wait((stdout_reader, stderr_reader), timeout=timeout)

    # Timeout doesn't raise an exception, but gevent.wait returns the greenlets
    # which did complete. So if both haven't completed, we kill them and fail
    # with a timeout.
    if len(greenlets) != 2:
        stdout_reader.kill()
        stderr_reader.kill()

        raise timeout_error()

    # Read the buffers into a list of lines
    stdout = stdout_reader.get()
    stderr = stderr_reader.get()

    logger.debug('--> Waiting for exit status...')
    process.wait()

    # Close any open file descriptor
    process.stdout.close()

    logger.debug('--> Command exit status: {0}'.format(process.returncode))
    return process.returncode == 0, stdout, stderr


def put_file(
    state, host, file_io, remote_file,
    sudo=False, sudo_user=None, su_user=None, print_output=False,
):
    with open(remote_file, 'w') as f:
        f.write(file_io.read())

@@ -262,16 +262,123 @@
 wargs):%0A
+    logger.info('%7B0%7D%7B1%7D'.format(%0A        host.print_prefix,%0A        click.style('Ready', 'green'),%0A    ))%0A%0A
     retu
@@ -569,24 +569,706 @@
 t=False,%0A):%0A
+    '''%0A    Execute a command on the local machine.%0A%0A    Args:%0A        state (%60%60pyinfra.api.State%60%60 obj): state object for this command%0A        hostname (string): hostname of the target%0A        command (string): actual command to execute%0A        sudo (boolean): whether to wrap the command with sudo%0A        sudo_user (string): user to sudo to%0A        get_pty (boolean): whether to get a PTY before executing the command%0A        env (dict): envrionment variables to set%0A        timeout (int): timeout for this command to complete before erroring%0A%0A    Returns:%0A        tuple: (exit_code, stdout, stderr)%0A        stdout and stderr are both lists of strings from each buffer.%0A    '''%0A%0A
     command 
@@ -3313,8 +3313,301 @@
 read())%0A
+%0A    if print_output:%0A        print('%7B0%7Dfile copied: %7B1%7D'.format(host.print_prefix, remote_file))%0A%0A%0Adef get_file(%0A    state, host, remote_file, file_io,%0A    sudo=False, sudo_user=False, su_user=None, print_output=False,%0A):%0A    with open(remote_file, 'r') as f:%0A        file_io.write(f.read())%0A

Add `put_file` in local connector and tidy up/docstring/comment.