import os
import posixpath
import sys
import threading
from collections import namedtuple

from parallels.core.logging import get_logger
from parallels.core import messages
from parallels.core.runners.base import BaseRunner
from parallels.core.runners.unix.ssh import SSHRunner
from parallels.core.utils import yaml_utils
from parallels.core.utils.common import is_empty, default, cached
from parallels.core.utils.common.threading_utils import synchronized_by_args

logger = get_logger(__name__)

KeyInfo = yaml_utils.pretty_yaml(
    namedtuple('KeyInfo', ('key_pathname', 'public_key'))
)


DEFAULT_AUTHORIZED_KEYS_LOCATIONS = [
    # On certain configurations, there is only the second file enabled,
    # so configure keys in both of the files
    '/root/.ssh/authorized_keys',
    '/root/.ssh/authorized_keys2'
]


class SSHKeyPool(object):
    def __init__(self, filename):
        self.keys = dict()
        self.lock = threading.Lock()
        self.filename = filename
        self.readonly = False

    @synchronized_by_args
    def get(self, source_server, target_server):
        """
        :return: parallels.core.utils.ssh_key_pool.KeyInfo
        """
        with self.lock:
            if self.readonly:
                raise Exception(messages.SSH_KEY_POOL_IS_READ_ONLY)

            if target_server.is_local() and source_server.ssh_key() is not None:
                return KeyInfo(
                    key_pathname=source_server.ssh_key(),
                    # we have no information about public key, we don't require to specify it in configuration file
                    # and actually it is not necessary
                    public_key=None
                )

            if (source_server, target_server,) not in self.keys:
                logger.debug(
                    messages.DEBUG_CONFIGURE_SSH_KEY_AUTH,
                    target_server.description(), source_server.description()
                )
                key_pathname, public_key = self._set_up_keys(
                    source_server, target_server,
                    authorized_keys_paths=self._get_server_authorized_keys_paths(source_server)
                )
                self.keys[(source_server, target_server)] = KeyInfo(
                    key_pathname, public_key
                )
                self._save_keys()

            return self.keys[(source_server, target_server)]

    def remove_all(self):
        """
        :rtype: None
        """
        for (source_server, target_server), key_info in self.keys.items():
            logger.debug(
                messages.DEBUG_REMOVE_SSH_KEY_AUTH,
                target_server.description(),
                source_server.description()
            )
            try:
                self._remove_keys(
                    source_server, target_server,
                    key_pathname=key_info.key_pathname,
                    public_key=key_info.public_key,
                    authorized_keys_paths=self._get_server_authorized_keys_paths(source_server)
                )
            except Exception as e:
                logger.exception()
                logger.error(
                    messages.FAILED_TO_REMOVE_SSH_KEY,
                    target_server.description(), source_server.description(), e
                )
        self.keys = {}
        self._save_keys()
        self.readonly = True

    def has_keys(self):
        """
        :rtype: bool
        """
        return len(self.keys) > 0

    def _save_keys(self):
        keys = {}
        for (source_server, target_server), key_info in self.keys.items():
            keys[source_server.ip(), target_server.ip()] = key_info
        yaml_utils.write_yaml(self.filename, keys)

    @staticmethod
    def _get_server_authorized_keys_paths(server):
        if (
            hasattr(server, 'node_settings') and
            hasattr(server.node_settings, 'ssh_authorized_keys_file') and
            not is_empty(server.node_settings.ssh_authorized_keys_file)
        ):
            return [server.node_settings.ssh_authorized_keys_file]
        else:
            root_home_directory = SSHKeyPool._get_home_directory(server)

            if sys.version_info.major >= 3:
                root_home_directory = root_home_directory.decode('utf-8')

            return [
                # On certain configurations, there is only the second file enabled,
                # so configure keys in both of the files
                posixpath.join(root_home_directory, '.ssh/authorized_keys'),
                posixpath.join(root_home_directory, '.ssh/authorized_keys2')
            ]

    @staticmethod
    @cached
    def _get_home_directory(server):
        with server.runner() as runner:
            assert isinstance(runner, BaseRunner)
            return runner.execute_command("echo ~").stdout.strip()

    @staticmethod
    def _set_up_keys(source_server, target_server, authorized_keys_paths=None):
        authorized_keys_paths = default(authorized_keys_paths, DEFAULT_AUTHORIZED_KEYS_LOCATIONS)
        target_ssh_keys_dir = target_server.get_session_file_path("ssh-keys")
        with target_server.runner() as runner:
            assert isinstance(runner, BaseRunner)
            runner.mkdir(target_ssh_keys_dir)
            private_key_filename = os.path.join(target_ssh_keys_dir, 'id_rsa.%s' % source_server.ip())
            public_key_filename = '%s.pub' % private_key_filename
            runner.execute_command(
                'echo y | ssh-keygen -q -t rsa -P "" -f {private_key_filename}',
                dict(
                    private_key_filename=private_key_filename
                )
            )
            public_key_contents = runner.get_file_contents(public_key_filename).strip()
            if sys.version_info.major >= 3:
                public_key_contents = public_key_contents.decode('utf-8')

        for authorized_keys_path in authorized_keys_paths:
            with source_server.runner() as runner:
                assert isinstance(runner, BaseRunner)
                if sys.version_info.major >= 3:
                    authorized_keys_path = str(authorized_keys_path)
                runner.mkdir(posixpath.dirname(authorized_keys_path))
                # Don't use upload_file_content function of runner, as rsync requires keys
                # we're trying to upload here. Use shell echo instead.
                runner.execute_command('echo {public_key_contents} >> {authorized_keys_path}', dict(
                    public_key_contents=public_key_contents, authorized_keys_path=authorized_keys_path
                ))

        return private_key_filename, public_key_contents

    @staticmethod
    def _remove_keys(source_server, target_server, key_pathname, public_key, authorized_keys_paths=None):
        authorized_keys_paths = default(authorized_keys_paths, DEFAULT_AUTHORIZED_KEYS_LOCATIONS)

        for authorized_keys_path in authorized_keys_paths:
            try:
                with source_server.runner() as runner:
                    assert isinstance(runner, BaseRunner)

                    remove_key_with_sed = True
                    if isinstance(runner, SSHRunner):
                        remove_key_with_sed = runner.is_utility_sed_available()

                    if remove_key_with_sed:
                        sed_expr = '\|{public_key}|d'.format(public_key=public_key)
                        runner.execute_command(
                            "sed -i -e {sed_expr} {authorized_keys_path}", dict(
                                sed_expr=sed_expr,
                                authorized_keys_path=authorized_keys_path
                            )
                        )
                    else:
                        initial_content = runner.get_file_contents(authorized_keys_path)
                        fixed_content = "\n".join(
                            line for line in initial_content.split("\n")
                            if public_key not in line
                        )
                        # do not use upload_file_content function as it relies on SSH keys presence
                        runner.execute_command(
                            'echo {content} > {path}', dict(path=authorized_keys_path, content=fixed_content)
                        )
            except Exception as e:
                logger.exception()
                logger.ferror(
                    messages.FAILED_TO_REMOVE_SSH_KEY_FROM_SOURCE,
                    filename=authorized_keys_path,
                    target=target_server.description(), source=source_server.description(), reason=e
                )

        try:
            with target_server.runner() as runner:
                runner.remove_file(key_pathname)
                runner.remove_file("%s.pub" % key_pathname)
        except Exception as e:
            logger.exception()
            logger.ferror(
                messages.FAILED_TO_REMOVE_SSH_KEY_FROM_TARGET,
                target=target_server.description(), source=source_server.description(), reason=e
            )
