import pipes
import sys
from contextlib import closing
import paramiko
from paramiko.ssh_exception import SSHException, AuthenticationException
from parallels.core.logging import get_logger
from parallels.core import messages, safe_utf8_encode
from parallels.core.utils.steps_profiler import sleep
from parallels.core import MigrationError
from parallels.core.utils.common import poll_data, safe_format

logger = get_logger(__name__)

POLL_INTERVALS = [5, 5, 5]


class ExecutionError(Exception):
    def __init__(self, command, exit_status, stdout, stderr):
        msg = messages.SSH_EXECUTION_ERROR % (command, exit_status, stdout, stderr)
        Exception.__init__(self, msg)
        self.command = command
        self.exit_status = exit_status
        self.stdout = stdout
        self.stderr = stderr


# Use this to avoid missing keys warning at WARNING logging level (which makes migrator output a bit messy).
# Please note that "how should we handle missing keys" is a disputable question,
# and probably ignoring them is not the best way.
class IgnoreMissingKeyPolicy(paramiko.MissingHostKeyPolicy):
    def missing_host_key(self, client, hostname, key):
        pass


def run_unchecked(
    ssh_client, command, stdin_content=None, output_codepage='utf-8', error_policy='strict', env=None
):
    """
    command - is either an unicode string (u''), or byte string that consists of only ascii symbols
    """
    if sys.version_info.major >= 3:
        pass
    else:
        command = safe_utf8_encode(command)

    if env is not None:
        command = '%s%s' % (get_env_str(env), command)

    def exec_command():
        try:
            return ssh_client.exec_command(command)
        except SSHException as e:
            expected_messages = [u'Unable to open channel.', u'Channel closed.']
            if e.message in expected_messages:
                logger.exception()
                return None
            else:
                raise

    result = poll_data(exec_command, POLL_INTERVALS)
    if result is None:
        assert(ssh_client.hostname is not None)
        raise MigrationError(messages.UNABLE_OPEN_CHANNEL_SSH_CONNECTION_S % ssh_client.hostname)
    else:
        stdin, stdout, stderr = result

    channel = stdout.channel

    if stdin_content is not None:
        stdin.write(stdin_content)

    # Close stdin so program waiting for input will get EOF.
    # stdin.close() does nothing beside flushing and setting internal 'closed' flag,
    # so it is needed to explicitly close writing side of channel.
    stdin.close()
    channel.shutdown_write()

    # TODO: This will hang if tool produces a lot of stderr output.
    stdout_content = stdout.read()
    if output_codepage is not None:
        stdout_content = stdout_content.decode(output_codepage, error_policy)
    stderr_content = stderr.read()
    if output_codepage is not None:
        stderr_content = stderr_content.decode(output_codepage, error_policy)

    stdout.close()
    stderr.close()

    exit_status = channel.recv_exit_status()

    return exit_status, stdout_content, stderr_content


class SSHClientWithHostname(paramiko.SSHClient):
    def __init__(self, max_connect_attempts=5, interval_between_connect_attempts=10):
        paramiko.SSHClient.__init__(self)
        self.hostname = None
        self._max_connect_attempts = max_connect_attempts
        self._interval_between_connect_attempts = interval_between_connect_attempts

    def connect(self, hostname, *args, **kw):
        self.hostname = hostname
        self.args = args
        self.kw = kw
        self._connect_multiple_attemtps()

    def _connect_multiple_attemtps(self):
        use_disabled_algorithms = False
        for attempt in range(0, self._max_connect_attempts):
            try:
                paramiko.SSHClient.connect(
                    self,
                    self.hostname,
                    disabled_algorithms={'pubkeys': ["rsa-sha2-512", "rsa-sha2-256"]} if use_disabled_algorithms else {},
                    *self.args,
                    **self.kw
                )
                if attempt > 0:
                    logger.info(messages.SUCCESSFULLY_CONNECTED_S_BY_SSH, self.hostname)
                return
            except (IOError, SSHException, AuthenticationException) as e:
                logger.exception()
                if attempt >= self._max_connect_attempts - (0 if use_disabled_algorithms else 1):
                    raise MigrationError(
                        safe_format(
                            messages.UNABLE_CONNECT_HOST_BY_SSH_EXCEPTION,
                            host=self.hostname, exception=e
                        )
                    )
                else:
                    if isinstance(e, SSHException) and not use_disabled_algorithms:
                        use_disabled_algorithms = True
                        logger.debug(
                            safe_format(
                                messages.SSH_KEYS_EXCEPTION_OCCURRED,
                                host=self.hostname,
                                exception=e
                            )
                        )
                    elif isinstance(e, AuthenticationException):
                        raise MigrationError(
                            safe_format(
                                messages.UNABLE_CONNECT_HOST_BY_SSH_EXCEPTION,
                                host=self.hostname, exception=e
                            )
                        )
                    else:
                        logger.error(
                            safe_format(
                                messages.SSH_UNABLE_TO_CONNECT_TO_HOST,
                                host=self.hostname,
                                interval_between_attempts=self._interval_between_connect_attempts,
                                exception=e
                            )
                        )
                        sleep(self._interval_between_connect_attempts, messages.WAIT_RECONNECT_BY_SSH)

    def exec_command(self, *args, **kwargs):
        if self.get_transport() is None or not self.get_transport().is_active():
            logger.warning(messages.SSH_CONNECTION_S_WAS_UNEXPECTEDLY_CLOSED % self.hostname)
            self.reconnect()

        return super(SSHClientWithHostname, self).exec_command(*args, **kwargs)

    def reconnect(self):
        super(SSHClientWithHostname, self).close()
        self._connect_multiple_attemtps()


def connect(settings, max_connect_attempts=5, interval_between_connect_attempts=10):
    client = SSHClientWithHostname(max_connect_attempts, interval_between_connect_attempts)
    client.set_missing_host_key_policy(IgnoreMissingKeyPolicy())
    settings.ssh_auth.connect(settings.ip, client)
    return closing(client)


def get_env_str(env, is_win=False):
    """Transform env variables dict into one line string which can be used to prefix command execution.

    Linux:
        dict(ENV_VAR_1=12, ENV_VAR_2=test) -> ENV_VAR_1=12 ENV_VAR_2=test

    Windows:
        dict(ENV_VAR_1=12, ENV_VAR_2=test) -> SET ENV_VAR_1=12 && SET ENV_VAR_2=test &&

    :type env: dict | None
    :type is_win: bool
    :rtype: unicode
    """
    if env:
        if sys.version_info.major >= 3:
            return '%s ' % ' '.join(
                "%s=%s" % (name, pipes.quote(value)) for name, value in env.items()
            )
        else:
            return '%s ' % ' '.join(
                "%s=%s" % (name, pipes.quote(safe_utf8_encode(value))) for name, value in env.items()
            )
    else:
        return ''
