"""
Remote controlling of nodes: copying files, executing commands

Copyright (c) 2010-2011 Mika Eloranta
See LICENSE for details.

"""

from __future__ import with_statement

import os
import subprocess
import logging
import shutil
import sys
import select
import errno
from . import errors
from . import colors


DONE = 0
STDOUT = 1
STDERR = 2


class RemoteControl:
    def __init__(self, node):
        self.node = node
        self.warn_timeout = 30.0 # seconds to wait before warning user after receiving any output
        self.terminate_timeout = 300.0 # seconds to wait before disconnecting after receiving any output

    def tag_line(self, tag, command, result=None, verbose=False, color=None):
        assert color
        if not verbose:
            return

        desc = "%s (%s): %s" % (color(self.node.name, "node"),
                                color(self.node.get("host"), "host"),
                                color(command, "command"))
        if result is not None:
            tag = "%s (%s)" % (tag, result)

        out_line = "%s %s %s %s\n" % (color("---", "header"),
                                      tag,
                                      desc,
                                      color("---", "header"))
        sys.stdout.write(out_line)

    def get_color(self, color, out_file=None):
        out_file = out_file or sys.stdout
        if color:
            return color
        else:
            return colors.Output(out_file, color="no").color

    def execute(self, command, verbose=False, color=None, output_lines=None, output_file=None, quiet=False):
        if output_file is not None:
            stdout_file = output_file
            stderr_file = output_file
        elif not quiet:
            stdout_file = sys.stdout
            stderr_file = sys.stderr
        else:
            stdout_file = None
            stderr_file = None

        result = None
        output_chunks = []
        color = self.get_color(color, out_file=stdout_file)
        self.tag_line(color("BEGIN", "header"), command, verbose=verbose,
                      color=color)

        try:
            while True:
                for code, output in self.execute_command(command):
                    if code == STDOUT:
                        if output_lines is not None:
                            output_chunks.append(output)
                        elif stdout_file:
                            if not verbose or stdout_file:
                                stdout_file.write(output)
                            else:
                                for line in output.splitlines(True):
                                    stdout_file.write(
                                        "[%s] %s" % (color(self.node.name,
                                                           "node"), line))
                            stdout_file.flush()
                    elif code == STDERR:
                        if not verbose or stdout_file:
                            stderr_file.write(output)
                        elif stderr_file:
                            for line in output.splitlines(True):
                                stderr_file.write(
                                    "{%s} %s" % (color(self.node.name,
                                                       "node"), line))
                        stderr_file.flush()
                    else: # DONE
                        if output_lines is not None:
                            output_lines.extend(
                                ("".join(output_chunks)).splitlines())

                        if not output:
                            result = color("OK", "op_ok")
                        else:
                            result = color(output, "op_error")
                        return output
        except Exception, error:
            result = color("%s: %s" % (error.__class__.__name__, error),
                           "op_error")
            raise
        finally:
            self.tag_line(color("END", "header"), command, result=result,
                          verbose=verbose, color=color)

    def shell(self, verbose=False, color=None):
        color = self.get_color(color)
        self.tag_line("BEGIN", "shell", verbose=verbose, color=color)
        try:
            return self.execute_shell()
        finally:
            self.tag_line("END", "shell", verbose=verbose, color=color)

    def close(self):
        pass

    def stat(self, file_path):
        assert 0, "must implement in sub-class"

    def read_file(self, file_path):
        assert 0, "must implement in sub-class"

    def put_file(self, source_path, dest_path, callback=None):
        assert 0, "must implement in sub-class"

    def write_file(self, file_path, contents, mode=None):
        assert 0, "must implement in sub-class"

    def execute_command(self, command):
        assert 0, "must implement in sub-class"

    def execute_shell(self):
        assert 0, "must implement in sub-class"

    def makedirs(self, dir_path):
        assert 0, "must implement in sub-class"


def convert_local_errors(method):
    """Convert local file-access errors to errors.RemoteError"""
    def wrapper(self, *args, **kw):
        try:
            return method(self, *args, **kw)
        except IOError, error:
            if error.errno == errno.ENOENT:
                raise errors.RemoteFileDoesNotExist(str(error))
            else:
                raise errors.RemoteError("%s: %s" % (error.__class__.__name__,
                                                     error))
        except OSError, error:
            raise errors.RemoteError("%s: %s" % (error.__class__.__name__,
                                                 error))

    wrapper.__doc__ = method.__doc__
    wrapper.__name__ = method.__name__

    return wrapper


class LocalControl(RemoteControl):
    """Local file-system access"""
    def __init__(self, node):
        RemoteControl.__init__(self, node)

    @convert_local_errors
    def put_file(self, source_path, dest_path, callback=None):
        shutil.copy(source_path, dest_path)

    @convert_local_errors
    def makedirs(self, dir_path):
        return os.makedirs(dir_path)

    @convert_local_errors
    def read_file(self, file_path):
        return file(file_path, "rb").read()

    @convert_local_errors
    def write_file(self, file_path, contents, mode=None):
        f = file(file_path, "wb")
        if mode is not None:
            os.chmod(file_path, mode)

        f.write(contents)
        f.close()

    @convert_local_errors
    def execute_command(self, cmd):
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)
        fds = [process.stdout, process.stderr]
        CHUNK = 2**20
        while process.poll() is None:
            r, w, e = select.select(fds, [], fds)
            if process.stdout in r:
                yield STDOUT, process.stdout.read(CHUNK)

            if process.stderr in r:
                yield STDERR, process.stderr.read(CHUNK)

        yield DONE, process.returncode

    @convert_local_errors
    def execute_shell(self):
        shell = os.environ.get("SHELL")
        if not shell:
            raise errors.RemoteError("$SHELL is not defined")

        return self.execute(["/usr/bin/env", shell])

    @convert_local_errors
    def stat(self, file_path):
        return os.stat(file_path)

    @convert_local_errors
    def utime(self, file_path, times):
        os.utime(file_path, times)


class SshRemoteControl(RemoteControl):
    """Remote control over an SSH connection using Paramiko"""
    def __init__(self, node):
        self.log = logging.getLogger("ssh")
        RemoteControl.__init__(self, node)
        self.key_filename = None
        cloud_key = node.get("cloud", {}).get("key-pair")
        if cloud_key:
            # cloud 'key-pair' overrides 'ssh-key' from host properties
            self.key_filename = "%s.pem" % cloud_key
        else:
            self.key_filename = node.get_tree_property("ssh-key")

        self.connect_timeout = node.get_tree_property("ssh-timeout", 60.0)


