import os
import pipes
import sys

from parallels.core import messages, MigrationNoContextError, MigrationError
from parallels.core.logging import get_logger
from parallels.core.registry import Registry
from parallels.core.runners.base import BaseRunner
from parallels.core.runners.entities import ExecutionResults, ExecutionOptions
from parallels.core.runners.exceptions.non_zero_exit_code import NonZeroExitCodeException
from parallels.core.runners.unix.base import UnixRunner
from parallels.core.utils import unix_utils, ssh_utils
from parallels.core.utils.common import default, is_run_on_windows, safe_format, cached
from parallels.core.utils.common.ip import is_ipv6
from parallels.core.utils.common.threading_utils import synchronized
from parallels.core.utils.steps_profiler import get_default_steps_profiler
from parallels.core.utils.unix_utils import format_command

logger = get_logger(__name__)
profiler = get_default_steps_profiler()


class SSHRunner(UnixRunner):
    LIST_FILES_MODE_LS = 'list-files-mode-ls'
    LIST_FILES_MODE_FIND = 'list-files-mode-find'
    LIST_FILES_MODE_AUTO = 'list-files-mode-auto'
    COPY_FILES_MODE_RSYNC = 'copy-files-mode-rsync'
    COPY_FILES_MODE_SCP = 'copy-files-mode-scp'
    COPY_FILES_MODE_AUTO = 'copy-files-mode-auto'

    def __init__(self, ssh, server, list_files_mode=LIST_FILES_MODE_FIND, copy_files_mode=COPY_FILES_MODE_RSYNC):
        super(SSHRunner, self).__init__(host_description=server.description(), hostname=server.ip())
        self.ssh = ssh
        self._server = server
        self._list_files_mode = list_files_mode
        self._copy_files_mode = copy_files_mode

        try:
            if self.execute_command('/bin/true').stdout != b'':
                raise TestCommandGarbageOutputError(safe_format(
                    messages.SSH_GARBAGE_OUTPUT, server_description=server.description(), server_ip=server.ip()
                ))
        except NonZeroExitCodeException:
            logger.exception()
            raise TestCommandNonZeroExitCodeFailure(safe_format(
                messages.SSH_NON_ZERO_EXIT_CODE, server_description=server.description(), server_ip=server.ip()
            ))

    def _execute_command_raw(self, command, arguments, execution_options):
        """Execute command and return results, but do not log, do not check for exit code, and so on.

        :type command: str | unicode 
        :type arguments: dict[str | unicode, str | unicode] 
        :type execution_options: parallels.core.runners.entities.ExecutionOptions 
        :rtype: parallels.core.runners.entities.ExecutionResults 
        """
        if execution_options.redirect_output_file is not None:
            raise NotImplementedError()

        command_str = self._format_command(command, arguments)
        if execution_options.working_dir is not None:
            command_str = "cd {working_dir}; {command}".format(
                working_dir=pipes.quote(execution_options.working_dir), command=command_str
            )
        env = self._get_env(execution_options.environment_variables)

        exit_code, stdout, stderr = ssh_utils.run_unchecked(
            self.ssh, command_str, execution_options.stdin_content,
            output_codepage=None, error_policy=None, env=env
        )

        return ExecutionResults(exit_code, stdout, stderr)

    @staticmethod
    def _get_env(env):
        env = default(env, {})
        # Execute all commands under 'en_US.utf-8' locale, if it is not specified,
        # so migrator works in the same way regardless of server's locale
        if 'LANG' not in env:
            env['LANG'] = 'en_US.utf-8'
        return env

    def upload_file(self, local_filename, remote_filename):
        """
        :type local_filename: str | unicode
        :type remote_filename: str | unicode
        :rtype: None
        """
        self._transfer_file(remote_filename, local_filename, direction='from_local_to_remote')

    @synchronized
    def upload_file_content(self, filename, content):
        global_context = Registry.get_instance().get_context()
        local_temp_file = global_context.migrator_server.get_session_file_path(
            global_context.local_temp_filename.get('ssh_upload')
        )
        with open(local_temp_file, 'wb') as fp:
            fp.write(content)

        try:
            self._transfer_file(filename, local_temp_file, direction='from_local_to_remote')
        finally:
            with global_context.migrator_server.runner() as runner:
                runner.remove_file(local_temp_file)

    def get_file(self, remote_filename, local_filename):
        self._transfer_file(remote_filename, local_filename, direction='from_remote_to_local')
        if not is_run_on_windows():
            import grp
            try:
                frontend_group_entry = grp.getgrnam('psaadm')
            except KeyError:
                frontend_group_entry = None

            if frontend_group_entry is not None:
                frontend_group_id = frontend_group_entry.gr_gid
                os.chown(local_filename, -1, frontend_group_id)

    def _transfer_file(self, remote_filename, local_filename, direction):
        """
        :type remote_filename: str | unicode
        :type local_filename: str | unicode
        :type direction: str | unicode
        :rtype: None
        """
        if self._copy_files_mode == self.COPY_FILES_MODE_RSYNC:
            self._transfer_file_with_rsync(remote_filename, local_filename, direction)
        elif self._copy_files_mode == self.COPY_FILES_MODE_SCP:
            self._transfer_file_with_scp(remote_filename, local_filename, direction)
        else: # if self._copy_files_mode == self.COPY_FILES_MODE_AUTO
            if self.is_utility_rsync_available():
                self._transfer_file_with_rsync(remote_filename, local_filename, direction)
            elif self.is_utility_scp_available():
                self._transfer_file_with_scp(remote_filename, local_filename, direction)
            else:
                raise MigrationError(safe_format(
                    messages.NO_RSYNC_OR_SCP, server_description=self._server.description()
                ))

    def _transfer_file_with_rsync(self, remote_filename, local_filename, direction):
        """
        :type remote_filename: str | unicode
        :type local_filename: str | unicode
        :type direction: str | unicode
        :rtype: None
        """
        global_context = Registry.get_instance().get_context()
        key_pathname = Registry.get_instance().get_context().ssh_key_pool.get(
            self._server, global_context.migrator_server
        ).key_pathname

        files = ["{ssh_user}@{remote_ip}:{remote_filename}", "{local_filename}"]
        if direction == 'from_local_to_remote':
            permissions = 'Fu=rw,Du=rwx,go='
            files.reverse()
        elif direction == 'from_remote_to_local':
            permissions = 'Fug=rw,Dug=rwx,o='
        else:
            raise Exception("Invalid direction passed to _transfer_file_with_rsync function")

        with global_context.migrator_server.runner() as runner:
            assert isinstance(runner, BaseRunner)
            port = self._server.settings().ssh_auth.port
            if sys.version_info.major >= 3:
                port = str(port)
            ssh_command = format_command(
                'ssh -i {key} -p {port} '
                '-o PasswordAuthentication=no -o StrictHostKeyChecking=no -o GSSAPIAuthentication=no',
                key=key_pathname,
                port=port,
            )
            runner.execute_command(
                "rsync -l --chmod={permissions} --timeout=30 -e {ssh_command} " + " ".join(files),
                dict(
                    local_filename=local_filename,
                    ssh_user=self._server.settings().ssh_auth.username,
                    remote_ip=self._escape_ipv6_address(self._server.ip()),
                    remote_filename=remote_filename,
                    ssh_command=ssh_command,
                    permissions=permissions
                )
            )

    def _transfer_file_with_scp(self, remote_filename, local_filename, direction):
        """
        :type remote_filename: str | unicode
        :type local_filename: str | unicode
        :type direction: str | unicode
        :rtype: None
        """
        global_context = Registry.get_instance().get_context()
        key_pathname = Registry.get_instance().get_context().ssh_key_pool.get(
            self._server, global_context.migrator_server
        ).key_pathname

        files = ["{ssh_user}@{remote_ip}:{remote_filename}", "{local_filename}"]
        if direction == 'from_local_to_remote':
            files.reverse()
        elif direction == 'from_remote_to_local':
            pass
        else:
            raise Exception("Invalid direction passed to _transfer_file_with_scp function")

        with global_context.migrator_server.runner() as runner:
            assert isinstance(runner, BaseRunner)
            runner.execute_command(
                (
                    "scp -i {key} -P {port} " 
                    "-o PasswordAuthentication=no -o StrictHostKeyChecking=no -o GSSAPIAuthentication=no "
                    +
                    " ".join(files)
                ),
                dict(
                    key=key_pathname,
                    port=self._server.settings().ssh_auth.port,
                    local_filename=local_filename,
                    ssh_user=self._server.settings().ssh_auth.username,
                    remote_ip=self._escape_ipv6_address(self._server.ip()),
                    remote_filename=remote_filename,
                )
            )

    def get_file_contents(self, remote_filename):
        """Get contents of the file.

        :type remote_filename: str | unicode
        :rtype: str | unicode
        """
        logger.fdebug(
            messages.DEBUG_GET_CONTENTS_OF_FILE_WITH_CAT, filename=remote_filename, server=self.host_description
        )
        return self.execute_command(
            'cat {path}', dict(path=remote_filename),
            ExecutionOptions(
                log_output=False, log_execution=False
            )
        ).stdout

    def download_directory(self, remote_directory, local_directory):
        """Download directory from remote server to local

        :type remote_directory: str | unicode
        :type local_directory: str | unicode
        :rtype: None
        """
        self._transfer_directory_with_rsync(remote_directory, local_directory, direction='from_remote_to_local')
        if not is_run_on_windows():
            import grp
            try:
                frontend_group_entry = grp.getgrnam('psaadm')
            except KeyError:
                frontend_group_entry = None

            if frontend_group_entry is not None:
                frontend_group_id = frontend_group_entry.gr_gid
                os.chown(local_directory, -1, frontend_group_id)
                for root, dirs, files in os.walk(local_directory):
                    for d in dirs:
                        os.chown(os.path.join(root, d), -1, frontend_group_id)
                    for f in files:
                        os.chown(os.path.join(root, f), -1, frontend_group_id)

    def upload_directory(self, local_path, remote_server, remote_path):
        """Upload directory from local server to remote

        :type local_path: str | unicode
        :type remote_server: parallels.core.connections.server.Server
        :type remote_path: str | unicode
        :rtype: None
        """
        self._transfer_directory_with_rsync(remote_path, local_path, direction='from_local_to_remote')

    def _transfer_directory_with_rsync(self, remote_directory, local_directory, direction):
        """Transfer directory between server using rsync

        :type remote_directory: str | unicode
        :type local_directory: str | unicode
        :type direction: str | unicode
        :rtype: None
        """
        global_context = Registry.get_instance().get_context()
        key_pathname = Registry.get_instance().get_context().ssh_key_pool.get(
            self._server, global_context.migrator_server
        ).key_pathname

        directories = ["{ssh_user}@{remote_ip}:{remote_directory}", "{local_directory}"]
        if direction == 'from_local_to_remote':
            permissions = 'Fu=rw,Du=rwx,go='
            directories.reverse()
        elif direction == 'from_remote_to_local':
            permissions = 'Fug=rw,Dug=rwx,o='
        else:
            raise Exception("Invalid direction passed to _transfer_directory_with_rsync function")

        with global_context.migrator_server.runner() as runner:
            assert isinstance(runner, BaseRunner)
            ssh_command = format_command(
                'ssh -i {key} -p {port} '
                '-o PasswordAuthentication=no -o StrictHostKeyChecking=no -o GSSAPIAuthentication=no',
                key=key_pathname,
                port=self._server.settings().ssh_auth.port,
            )
            runner.execute_command(
                "rsync -lr --chmod={permissions} --timeout=30 -e {ssh_command} " + " ".join(directories),
                dict(
                    local_directory=local_directory + '/',
                    ssh_user=self._server.settings().ssh_auth.username,
                    remote_ip=self._escape_ipv6_address(self._server.ip()),
                    remote_directory=remote_directory + '/',
                    ssh_command=ssh_command,
                    permissions=permissions
                )
            )

    def move(self, src_path, dst_path):
        self.execute_command(
            '/bin/mv {src_path} {dst_path}', dict(src_path=src_path, dst_path=dst_path)
        )

    def get_files_list(self, path):
        if self._list_files_mode == self.LIST_FILES_MODE_FIND:
            return unix_utils.get_files_list_with_find(self, path)
        elif self._list_files_mode == self.LIST_FILES_MODE_LS:
            return unix_utils.get_files_list_with_ls(self, path)
        else:  # if self._list_files_mode == self.LIST_FILES_MODE_AUTO:
            if self.is_utility_find_available():
                return unix_utils.get_files_list_with_find(self, path)
            elif self.is_utility_ls_available():
                return unix_utils.get_files_list_with_ls(self, path)
            else:
                raise MigrationError(safe_format(
                    messages.NO_LS_OR_FIND, server_description=self._server.description()
                ))

    def get_dir_tree(self, path):
        """Retrieve tree structure of the directory with given path, like the following:

        [
            'file_1',
            'file_2',
            {
                'dir_1': [
                    'file_3'
                ]
            }
        ]

        :type path: str | unicode
        :rtype: list[str | {str, list}]
        """
        if self._list_files_mode == self.LIST_FILES_MODE_FIND:
            # get directory tree with a single "find" command
            return unix_utils.get_dir_tree_with_find(self, path)
        elif self._list_files_mode == self.LIST_FILES_MODE_LS:
            # use default implementation - recursively traverse directory tree with "ls" command
            return super(SSHRunner, self).get_dir_tree(path)
        else:  # if self._list_files_mode == self.LIST_FILES_MODE_AUTO:
            if self.is_utility_find_available():
                # get directory tree with a single "find" command, when command is available
                return unix_utils.get_dir_tree_with_find(self, path)
            elif self.is_utility_ls_available():
                # use default implementation - recursively traverse directory tree with "ls" command when "find" is
                # not available
                return super(SSHRunner, self).get_dir_tree(path)
            else:
                raise MigrationError(safe_format(
                    messages.NO_LS_OR_FIND, server_description=self._server.description()
                ))

    def get_directories_list(self, path):
        """
        :type path: str | unicode
        :rtype: list[str | unicode]
        """
        if self._list_files_mode == self.LIST_FILES_MODE_FIND:
            return unix_utils.get_directories_list_with_find(self, path)
        elif self._list_files_mode == self.LIST_FILES_MODE_LS:
            return unix_utils.get_directories_list_with_ls(self, path)
        else:  # if self._list_files_mode == self.LIST_FILES_MODE_AUTO:
            if self.is_utility_find_available():
                return unix_utils.get_directories_list_with_find(self, path)
            elif self.is_utility_ls_available():
                return unix_utils.get_directories_list_with_ls(self, path)
            else:
                raise MigrationError(safe_format(
                    messages.NO_LS_OR_FIND, server_description=self._server.description()
                ))

    def iter_files_list_nested(self, path):
        """Iterate over all (including files in nested directories) file names inside a path

        :type path: str | unicode
        :rtype: list[str | unicode]
        """
        logger.fdebug(
            messages.DEBUG_LIST_FILES_OF_DIRECTORY_WITH_FIND, directory=path, server=self.host_description
        )
        return self.execute_command(
            'find {path} -print0', dict(path=path), ExecutionOptions(log_execution=False)
        ).stdout.split("\0")

    def get_file_size(self, remote_filename):
        """Get size of specified file (in bytes)

        :type remote_filename: str | unicode
        :rtype: int
        """
        logger.fdebug(
            messages.DEBUG_GET_SIZE_OF_FILE_WITH_STAT, filename=remote_filename, server=self.host_description
        )
        return int(
            self.execute_command(
                'stat --printf="%s" {file}', dict(file=remote_filename), ExecutionOptions(log_execution=False)
            ).stdout
        )

    def remove_file(self, filename):
        self.execute_command(
            '/bin/rm -f {filename}', dict(filename=filename)
        )

    def remove_directory(self, directory, is_remove_root=True):
        if not is_remove_root:
            raise NotImplementedError()

        self.execute_command('/bin/rm -rf {directory}', dict(directory=directory))

    def mkdir(self, dirname):
        """
        :type dirname: str | unicode
        :rtype: None
        """
        self.execute_command('/bin/mkdir -p {dirname}', dict(dirname=dirname))

    def file_exists(self, filename):
        return self.execute_command(
            'test -e {filename}', dict(filename=filename),
            ExecutionOptions(ignore_exit_code=True)
        ).exit_code == 0

    def is_dir(self, path):
        """Check whether provided path is a directory

        :type path: str | unicode
        :rtype: bool
        """
        return self.execute_command(
            'test -d {path}', dict(path=path),
            ExecutionOptions(ignore_exit_code=True)
        ).exit_code == 0

    def chmod(self, filename, mode):
        """Change file mode

        :type filename: str | unicode
        :type mode: int
        :rtype: None
        """
        self.execute_command(
            'chmod {mode} {filename}', dict(
                mode="%o" % mode,  # octal representation of the mode - as required by "chmod" utility
                filename=filename
            )
        )

    @cached
    def is_utility_find_available(self):
        """
        :rtype: bool
        """
        return self._is_utility_available('find')

    @cached
    def is_utility_rsync_available(self):
        """
        :rtype: bool
        """
        return self._is_utility_available('rsync')

    @cached
    def is_utility_scp_available(self):
        """
        :rtype: bool
        """
        return self._is_utility_available('scp')

    @cached
    def is_utility_ls_available(self):
        """
        :rtype: bool
        """
        return self._is_utility_available('ls')

    @cached
    def is_utility_sed_available(self):
        """
        :rtype: bool
        """
        return self._is_utility_available('sed')

    def _is_utility_available(self, utility):
        """Check whether specified utility is available on the server or not

        :type utility: str | unicode
        :rtype: bool
        """
        execution_results = self.execute_command(
            "command -v {utility}", dict(utility=utility), execution_options=ExecutionOptions(ignore_exit_code=True)
        )
        return execution_results.exit_code == 0 and len(execution_results.stdout) > 0

    @staticmethod
    def _escape_ipv6_address(address):
        """Escape IP address for rsync/scp utilities

        :type address: str | unicode
        :rtype: str | unicode
        """
        if is_ipv6(address):
            return '[%s]' % address
        else:
            return address


class TestCommandNonZeroExitCodeFailure(MigrationNoContextError):
    pass


class TestCommandGarbageOutputError(MigrationNoContextError):
    pass
