import importlib
import logging
import os
import time
from base64 import b64decode
from dataclasses import dataclass
from functools import partial
from pathlib import Path, PurePath, PurePosixPath
from stat import S_ISLNK
from typing import Callable, List, Optional
from urllib.parse import quote, urljoin

from loguru import logger

from flexget import plugin
from flexget.entry import Entry
from flexget.task import TaskAbort

# retry configuration constants
RETRY_INTERVAL_SEC: int = 15
RETRY_STEP_SEC: int = 5
HOST_KEY_TYPES: dict = {
    'ssh-rsa': 'RSAKey',
    'ssh-ed25519': 'Ed25519Key',
}

try:
    import paramiko
    import pysftp

    logging.getLogger("paramiko").setLevel(logging.ERROR)
except ImportError:
    pysftp = None

NodeHandler = Callable[[str], None]

logger = logger.bind(name='sftp_client')


def _set_authentication_patch(self, password, private_key, private_key_pass):
    """
    Patch pysftp.Connection._set_authentication to support additional
    key types
    """
    if password is None:
        # Use Private Key.
        if not private_key:
            # Try to use default key.
            if os.path.exists(os.path.expanduser('~/.ssh/id_rsa')):
                private_key = '~/.ssh/id_rsa'
            elif os.path.exists(os.path.expanduser('~/.ssh/id_dsa')):
                private_key = '~/.ssh/id_dsa'
            else:
                raise pysftp.exceptions.CredentialException("No password or key specified.")

        if isinstance(private_key, (paramiko.AgentKey, paramiko.RSAKey)):
            # use the paramiko agent or rsa key
            self._tconnect['pkey'] = private_key
        else:
            # isn't a paramiko AgentKey or RSAKey, try to build a
            # key from what we assume is a path to a key
            private_key_file = os.path.expanduser(private_key)
            for key in [paramiko.RSAKey, paramiko.DSSKey, paramiko.Ed25519Key, paramiko.ECDSAKey]:
                try:  # try all the keys
                    self._tconnect['pkey'] = key.from_private_key_file(
                        private_key_file, private_key_pass
                    )
                    return
                except paramiko.SSHException:  # if it fails, try dss
                    pass
            raise paramiko.SSHException(f'Unknown key type: {private_key}')


@dataclass
class HostKey:
    """
    Host key used to connect to a SFTP server if not defined in known_hosts.
    """

    key_type: str
    public_key: str


class SftpClient:
    def __init__(
        self,
        host: str,
        port: int,
        username: str,
        password: Optional[str] = None,
        private_key: Optional[str] = None,
        private_key_pass: Optional[str] = None,
        host_key: Optional[HostKey] = None,
        connection_tries: int = 3,
    ):
        if not pysftp:
            raise plugin.DependencyError(
                issued_by='sftp_client',
                missing='pysftp',
                message='sftp client requires the pysftp Python module.',
            )

        self.host: str = host
        self.port: int = port
        self.username: str = username
        self.password: Optional[str] = password
        self.private_key: Optional[str] = private_key
        self.private_key_pass: Optional[str] = private_key_pass
        self.host_key: Optional[HostKey] = host_key

        self.prefix: str = self._get_prefix()
        self._sftp: pysftp.Connection = self._connect(connection_tries)
        self._handler_builder: HandlerBuilder = HandlerBuilder(
            self._sftp, self.prefix, self.private_key, self.private_key_pass, self.host_key
        )

    def list_directories(
        self,
        directories: List[str],
        recursive: bool,
        get_size: bool,
        files_only: bool,
        dirs_only: bool,
    ) -> List[Entry]:
        """
        Build a list of entries from a provided list of directories on an SFTP server
        :param directories: list of directories to generate entries for
        :param recursive: boolean indicating whether to list recursively
        :param get_size: boolean indicating whether to compute size for each node (potentially slow for directories)
        :param files_only: boolean indicating whether to exclude directories
        :param dirs_only: boolean indicating whether to exclude files
        :return: a list of entries describing the contents of the provided directories
        """

        entries: List[Entry] = []

        dir_handler: NodeHandler = self._handler_builder.get_dir_handler(
            get_size, files_only, entries
        )
        file_handler: NodeHandler = self._handler_builder.get_file_handler(
            get_size, dirs_only, entries
        )
        unknown_handler: NodeHandler = self._handler_builder.get_unknown_handler()

        for directory in directories:
            try:
                # Always normalize the root path so it's not necessary to normalised
                # nodes as there are discovered, which means that symlinks will appear
                # in the entry paths raw rather than been resolved to their target.
                self._sftp.walktree(
                    self._sftp.normalize(directory),
                    file_handler,
                    dir_handler,
                    unknown_handler,
                    recursive,
                )
            except OSError as e:
                logger.warning('Failed to open {} ({})', directory, str(e))
                continue

        return entries

    def download(self, source: str, to: str, recursive: bool, delete_origin: bool) -> None:
        """
        Downloads the file specified in "source" to the destination specified in "to"
        :param source: path of the resource to download
        :param to: path of the directory to download to
        :param recursive: indicates whether to download the contents of "source" recursively
        :param delete_origin: indicates whether to delete the source resource upon download, is the source
                              is a symlink, only the symlink will be removed rather than it's target.
        """

        dir_handler: NodeHandler = self._handler_builder.get_null_handler()
        unknown_handler: NodeHandler = self._handler_builder.get_unknown_handler()

        parsed_path: PurePosixPath = PurePosixPath(source)

        if not self.path_exists(source):
            raise SftpError(f'Remote path does not exist: {source}')

        is_symlink: bool = self.is_link(source)
        if self.is_file(source):
            source_file: str = parsed_path.name
            source_dir: str = parsed_path.parent.as_posix()
            try:
                self._sftp.cwd(source_dir)
                self._download_file(to, delete_origin and not is_symlink, source_file)
            except Exception as e:
                raise SftpError(f'Failed to download file {source} ({e!s})')

            if delete_origin and is_symlink:
                self.remove_file(source)

        elif self.is_dir(source):
            base_path: str = parsed_path.joinpath('..').as_posix()
            dir_name: str = parsed_path.name
            handle_file: NodeHandler = partial(
                self._download_file, to, delete_origin and not is_symlink
            )

            try:
                self._sftp.cwd(base_path)
                self._sftp.walktree(dir_name, handle_file, dir_handler, unknown_handler, recursive)
            except Exception as e:
                raise SftpError(f'Failed to download directory {source} ({e!s})')

            if delete_origin:
                if self.is_link(source):
                    self.remove_file(source)
                else:
                    self.remove_dir(source)
        else:
            logger.warning('Skipping unknown file: {}', source)

    def upload(self, source: str, to: str) -> None:
        """
        Upload files or directories to an SFTP server
        :param source: file or directory to upload
        :param to: destination
        """
        if Path(source).is_dir():
            logger.verbose('Skipping directory {}', source)  # type: ignore
        else:
            self._upload_file(source, to)

    def remove_dir(self, path: str) -> None:
        """
        Remove a directory if it's empty
        :param path: directory to remove
        """
        if self._sftp.exists(path) and not self._sftp.listdir(path):
            logger.debug('Attempting to delete directory {}', path)
            try:
                self._sftp.rmdir(path)
            except Exception as e:
                logger.error('Failed to delete directory {} ({})', path, str(e))

    def remove_file(self, path: str) -> None:
        """
        Remove a file if it's empty
        :param path: file to remove
        """
        logger.debug('Deleting remote file {}', path)
        try:
            self._sftp.remove(path)
        except Exception as e:
            logger.error('Failed to delete file {} ({})', path, str(e))
            return

    def is_file(self, path: str) -> bool:
        """
        Check if the node at a given path is a file
        :param path: path to check
        :return: boolean indicating if the path is a file
        """
        return self._sftp.isfile(path)

    def is_dir(self, path: str) -> bool:
        """
        Check if the node at a given path is a directory
        :param path: path to check
        :return: boolean indicating if the path is a directory
        """
        return self._sftp.isdir(path)

    def is_link(self, path: str) -> bool:
        """
        Check if the node at a given path is a directory
        :param path: path to check
        :return: boolean indicating if the path is a directory
        """
        return S_ISLNK(self._sftp.sftp_client.lstat(path).st_mode)

    def path_exists(self, path: str) -> bool:
        """
        Check of a path exists

        :param path: Path to check
        :return: boolean indicating if the path exists
        """
        return self._sftp.lexists(path)

    def make_dirs(self, path: str) -> None:
        """
        Build directories

        :param path: path to build
        """

        if not self.path_exists(path):
            try:
                self._sftp.makedirs(path)
            except Exception as e:
                raise SftpError(f'Failed to create remote directory {path} ({e!s})')

    def close(self) -> None:
        """
        Close the sftp connection
        """
        self._sftp.close()

    def set_socket_timeout(self, socket_timeout_sec):
        """
        Sets the SFTP client socket timeout
        :param socket_timeout_sec: Socket timeout in seconds
        """
        self._sftp.timeout = socket_timeout_sec

    def _connect(self, connection_tries: int) -> 'pysftp.Connection':
        tries: int = connection_tries
        retry_interval: int = RETRY_INTERVAL_SEC

        logger.debug('Connecting to {}', self.host)

        sftp: Optional[pysftp.Connection] = None

        while not sftp:
            try:
                pysftp.Connection._set_authentication = _set_authentication_patch
                sftp = pysftp.Connection(
                    host=self.host,
                    username=self.username,
                    private_key=self.private_key,
                    password=self.password,
                    port=self.port,
                    private_key_pass=self.private_key_pass,
                    cnopts=self._get_cnopts(),
                )
                logger.verbose('Connected to {}', self.host)  # type: ignore
            except Exception as e:
                tries -= 1
                logger.debug('Caught exception: {}', e)
                if not tries:
                    raise TaskAbort(f'Failed to connect to {self.host}')
                else:
                    logger.debug('Caught exception: {}', e)
                    logger.warning(
                        'Failed to connect to {}; waiting {} seconds before retrying.',
                        self.host,
                        retry_interval,
                    )
                    time.sleep(retry_interval)
                    retry_interval += RETRY_STEP_SEC

        return sftp

    def _get_cnopts(self) -> Optional['pysftp.CnOpts']:
        if not self.host_key:
            return None
        KeyClass = getattr(
            importlib.import_module("paramiko"), HOST_KEY_TYPES[self.host_key.key_type]
        )
        key = KeyClass(data=b64decode(self.host_key.public_key))
        cnopts = pysftp.CnOpts()
        cnopts.hostkeys.add(self.host, self.host_key.key_type, key)
        return cnopts

    def _upload_file(self, source: str, to: str) -> None:
        if not Path(source).exists():
            logger.warning('File no longer exists:', source)
            return

        destination = self._get_upload_path(source, to)
        destination_url: str = urljoin(self.prefix, destination)

        if not self.path_exists(to):
            try:
                self.make_dirs(to)
            except Exception as e:
                raise SftpError(f'Failed to create remote directory {to} ({e!s})')

        if not self.is_dir(to):
            raise SftpError(f'Not a directory: {to}')

        try:
            self._put_file(source, destination)
            logger.verbose('Successfully uploaded {} to {}', source, destination_url)  # type: ignore
        except OSError:
            raise SftpError(f'Remote directory does not exist: {to}')
        except Exception as e:
            raise SftpError(f'Failed to upload {source} ({e!s})')

    def _download_file(self, destination: str, delete_origin: bool, source: str) -> None:
        destination_path: str = self._get_download_path(source, destination)
        destination_dir: str = Path(destination_path).parent.as_posix()

        if Path(destination_path).exists():
            logger.verbose(  # type: ignore
                'Skipping {} because destination file {} already exists.', source, destination_path
            )
            return

        Path(destination_dir).mkdir(parents=True, exist_ok=True)

        logger.verbose('Downloading file {} to {}', source, destination)  # type: ignore

        try:
            self._sftp.get(source, destination_path)
        except Exception as e:
            logger.error('Failed to download {} ({})', source, e)
            if Path(destination_path).exists():
                logger.debug('Removing partially downloaded file {}', destination_path)
                Path(destination_path).unlink()
            raise e

        if delete_origin:
            self.remove_file(source)

    def _put_file(self, source: str, destination: str) -> None:
        return self._sftp.put(source, destination)

    def _get_prefix(self) -> str:
        """
        Generate SFTP URL prefix
        """

        def get_login_string() -> str:
            if self.username and self.password:
                return f'{self.username}:{self.password}@'
            elif self.username:
                return f'{self.username}@'
            else:
                return ''

        def get_port_string() -> str:
            if self.port and self.port != 22:
                return f':{self.port}'
            else:
                return ''

        login_string = get_login_string()
        host = self.host
        port_string = get_port_string()

        return f'sftp://{login_string}{host}{port_string}/'

    @staticmethod
    def _get_download_path(path: str, destination: str) -> str:
        return PurePosixPath(destination).joinpath(Path(path)).as_posix()

    @staticmethod
    def _get_upload_path(source: str, to: str):
        basename: str = PurePath(source).name
        return PurePosixPath(to, basename).as_posix()


class SftpError(Exception):
    pass


class HandlerBuilder:
    """
    Class for building pysftp.Connection.walktree node handlers.

    :param sftp: A Connection object
    :param logger: a logger object
    :param url_prefix: SFTP URL prefix
    """

    def __init__(
        self,
        sftp: 'pysftp.Connection',
        url_prefix: str,
        private_key: Optional[str],
        private_key_pass: Optional[str],
        host_key: Optional[HostKey],
    ):
        self._sftp = sftp
        self._prefix = url_prefix
        self._private_key = private_key
        self._private_key_pass = private_key_pass
        self._host_key = host_key

    def get_file_handler(
        self, get_size: bool, dirs_only: bool, entry_accumulator: list
    ) -> NodeHandler:
        """
        Builds a file node handler suitable for use with pysftp.Connection.walktree

        :param get_size: boolean indicating whether to compute the for each file
        :param dirs_only: boolean indicating whether to skip files
        :param entry_accumulator: list to add entries to
        """
        return partial(
            Handlers.handle_file,
            self._sftp,
            self._prefix,
            get_size,
            dirs_only,
            self._private_key,
            self._private_key_pass,
            self._host_key,
            entry_accumulator,
        )

    def get_dir_handler(
        self, get_size: bool, files_only: bool, entry_accumulator: list
    ) -> NodeHandler:
        """
        Builds a file node handler suitable for use with pysftp.Connection.walktree

        :param get_size: boolean indicating whether to compute the for each file
        :param files_only: Boolean indicating whether to skip directories
        :param entry_accumulator: list to add entries to
        """
        return partial(
            Handlers.handle_directory,
            self._sftp,
            self._prefix,
            get_size,
            files_only,
            self._private_key,
            self._private_key_pass,
            self._host_key,
            entry_accumulator,
        )

    def get_unknown_handler(self) -> NodeHandler:
        """
        Builds an unknown node handler suitable for use with pysftp.Connection.walktree
        """
        return partial(Handlers.handle_unknown)

    def get_null_handler(self) -> NodeHandler:
        """
        Builds a noop node handler suitable for use with pysftp.Connection.walktree
        """
        return partial(Handlers.null_node_handler)


class Handlers:
    @classmethod
    def handle_file(
        cls,
        sftp: 'pysftp.Connection',
        prefix: str,
        get_size: bool,
        dirs_only: bool,
        private_key: Optional[str],
        private_key_pass: Optional[str],
        host_key: Optional[HostKey],
        entry_accumulator: List[Entry],
        path: str,
    ) -> None:
        """
        File node handler. Adds a file entry to entry_accumulator.

        :param sftp: A pysftp.Connection object
        :param logger: a logger object
        :param prefix: SFTP URL prefix
        :param get_size: boolean indicating whether to compute the size of each file
        :param dirs_only: boolean indicating whether to skip files
        :param private_key: private key path
        :param private_key_pass: private key password
        :param host_key: Host key for the remote server if not in known_hosts
        :param entry_accumulator: a list in which to store entries
        :param path: path to handle
        """
        if dirs_only:
            return

        size_handler = partial(cls._file_size, sftp)
        entry = cls._get_entry(
            sftp, prefix, size_handler, get_size, path, private_key, private_key_pass, host_key
        )
        entry_accumulator.append(entry)

    @classmethod
    def handle_directory(
        cls,
        sftp: 'pysftp.Connection',
        prefix: str,
        get_size: bool,
        files_only: bool,
        private_key: Optional[str],
        private_key_pass: Optional[str],
        host_key: Optional[HostKey],
        entry_accumulator: List[Entry],
        path: str,
    ) -> None:
        """
        Directory node handler. Adds a directory entry to entry_accumulator.

        :param sftp: A pysftp.Connection object
        :param logger: a logger object
        :param prefix: SFTP URL prefix
        :param get_size: boolean indicating whether to compute the size of each directory
        :param files_only: Boolean indicating whether to skip directories
        :param entry_accumulator: a list in which to store entries
        :param private_key: private key path
        :param private_key_pass: private key password
        :param host_key: Host key for the remote server if not in known_hosts
        :param path: path to handle
        """
        if files_only:
            return

        dir_size: Callable[[str], int] = partial(cls._dir_size, sftp)
        entry: Entry = cls._get_entry(
            sftp, prefix, dir_size, get_size, path, private_key, private_key_pass, host_key
        )
        entry_accumulator.append(entry)

    @staticmethod
    def handle_unknown(path: str) -> None:
        """
        Handler for unknown nodes; logs a warning.

        :param logger: a logger object
        :param path: path to handle
        """
        logger.warning('Skipping unknown file: {}', path)

    @staticmethod
    def null_node_handler(path: str) -> None:
        """
        Generic noop node handler

        :param logger: a logger object
        :param path: path to handle
        :return:
        """
        logger.debug('null handler called  for {}', path)

    @staticmethod
    def _get_entry(
        sftp: 'pysftp.Connection',
        prefix: str,
        size_handler: Callable[[str], int],
        get_size,
        path: str,
        private_key: Optional[str],
        private_key_pass: Optional[str],
        host_key: Optional[HostKey],
    ) -> Entry:
        url = urljoin(prefix, quote(path))
        title = PurePosixPath(path).name

        entry = Entry(title, url)

        if get_size:
            try:
                size = size_handler(path)
            except Exception as e:
                logger.warning('Failed to get size for {} ({})', path, e)
                size = -1
            entry['content_size'] = size

        entry['private_key'] = private_key
        entry['private_key_pass'] = private_key_pass
        if host_key:
            entry['host_key'] = {
                'key_type': host_key.key_type,
                'public_key': host_key.public_key,
            }

        return entry

    @classmethod
    def _dir_size(cls, sftp: 'pysftp.Connection', path: str) -> int:
        sizes: List[int] = []

        size_accumulator = partial(cls._accumulate_file_size, sftp, sizes)
        sftp.walktree(path, size_accumulator, size_accumulator, size_accumulator, True)

        return sum(sizes)

    @classmethod
    def _accumulate_file_size(
        cls, sftp: 'pysftp.Connection', size_accumulator: List[int], path: str
    ) -> None:
        size_accumulator.append(cls._file_size(sftp, path))

    @staticmethod
    def _file_size(sftp: 'pysftp.Connection', path: str) -> int:
        """
        Helper function to get the size of a file node
        """
        return sftp.lstat(path).st_size
