import errno
import logging
import paramiko
import socket
from os import path
from stat import S_ISDIR, S_ISREG

_KEY_TYPES = {
    "dsa": paramiko.DSSKey,
    "rsa": paramiko.RSAKey,
    "ecdsa": paramiko.ECDSAKey,
    "ed25519": paramiko.Ed25519Key,
}


class SFTPController(paramiko.SFTPClient):
    def __init__(self, sock):
        super().__init__(sock)

    def exists(self, path):
        try:
            self.stat(path)
        except IOError as e:
            return e.errno != errno.ENOENT
        return True

    def list_dirs(self, path):
        return [
            d.filename for d in self.listdir_attr(path) if S_ISDIR(d.st_mode)
        ]

    def list_files(self, path):
        return [
            f.filename for f in self.listdir_attr(path) if S_ISREG(f.st_mode)
        ]

    @classmethod
    def from_transport(cls, t):
        chan = t.open_session()
        chan.invoke_subsystem("sftp")
        return cls(chan)


class SSHController:
    def __init__(
        self,
        host,
        user,
        key_path=None,
        key_password=None,
        key_type="rsa",
        ssh_password=None,
        port=22,
    ):
        self.host = host
        self.user = user
        self.ssh_password = ssh_password if key_path is None else None
        self.port = port
        self.nb_bytes = 1024
        self.keys, self.transport = [], None
        key_type = key_type.lower()

        if key_path:
            self.keys.append(
                _KEY_TYPES[key_type].from_private_key(
                    open(path.expanduser(key_path), 'r'),
                    key_password,
                )
            )
        elif ssh_password is None:
            self.keys = paramiko.Agent().get_keys()

            try:
                key_file = _KEY_TYPES[key_type].from_private_key(
                    open(path.expanduser(f"~/.ssh/id_{key_type}"), 'r'),
                    key_password
                )
            except Exception:
                pass
            else:
                self.keys.insert(
                    len(self.keys) if key_password is None else 0, key_file
                )

            if not self.keys:
                logging.error("No valid key found")

    def connect(self):
        try:
            ssh_socket = socket.create_connection((self.host, self.port))
        except OSError as e:
            logging.error(f"Connection failed: {e.strerror}")
            return 1

        self.transport = paramiko.Transport(ssh_socket)

        if self.ssh_password is not None:
            try:
                self.transport.connect(
                    username=self.user,
                    password=self.ssh_password,
                )
            except paramiko.SSHException:
                pass
        else:
            for key in self.keys:
                try:
                    self.transport.connect(username=self.user, pkey=key)
                except paramiko.SSHException:
                    continue
                break

        if not self.transport.is_authenticated():
            logging.error("SSH negotiation failed")
            return 1

        logging.info(f"Successfully connected to {self.user}@{self.host}")
        return 0

    def __run_until_event(
        self,
        command,
        stop_event,
        display=True,
        combine_stderr=False,
        capture_output=False,
    ):
        channel = self.transport.open_session()
        output = ""
        timeout = 2

        channel.settimeout(timeout)
        channel.set_combine_stderr(combine_stderr)
        channel.get_pty()
        channel.exec_command(command)

        if not display and not capture_output:
            stop_event.wait()
        else:
            while True:
                try:
                    raw_data = channel.recv(self.nb_bytes)
                except socket.timeout:
                    if stop_event.is_set():
                        break
                    continue

                if not len(raw_data):
                    break

                data = raw_data.decode("utf-8")

                if display:
                    print(data, end='')

                if capture_output:
                    output += data

                if stop_event.is_set():
                    break

        channel.close()
        return (channel.exit_status_ready(), output.splitlines())

    def __run_until_exit(
        self,
        command,
        timeout,
        display=True,
        combine_stderr=False,
        capture_output=False,
    ):
        channel = self.transport.open_session()
        output = ""

        channel.settimeout(timeout)
        channel.set_combine_stderr(combine_stderr)
        channel.get_pty()
        channel.exec_command(command)

        try:
            if not display and not capture_output:
                return (channel.recv_exit_status(), output.splitlines())
            else:
                while True:
                    raw_data = channel.recv(self.nb_bytes)

                    if not len(raw_data):
                        break

                    data = raw_data.decode("utf-8")

                    if display:
                        print(data, end='')

                    if capture_output:
                        output += data
        except socket.timeout:
            logging.warning(f"Timeout after {timeout}s")
            return (1, output.splitlines())
        except KeyboardInterrupt:
            logging.info("KeyboardInterrupt")
            return (0, output.splitlines())
        finally:
            channel.close()

        return (channel.recv_exit_status(), output.splitlines())

    def run(
        self,
        command,
        display=False,
        combine_stderr=False,
        capture_output=False,
        stop_event=None,
        timeout=600,
    ):
        if stop_event:
            return self.__run_until_event(
                command,
                stop_event,
                display=display,
                combine_stderr=combine_stderr,
                capture_output=capture_output,
            )
        else:
            return self.__run_until_exit(
                command,
                timeout,
                display=display,
                combine_stderr=combine_stderr,
                capture_output=capture_output,
            )

    def disconnect(self):
        if self.transport:
            self.transport.close()

    def __getattr__(self, target):
        def wrapper(*args, **kwargs):
            if not self.transport.is_authenticated():
                logging.error("SSH session is not ready")
                return 1
            sftp_channel = SFTPController.from_transport(self.transport)
            r = getattr(sftp_channel, target)(*args, **kwargs)
            sftp_channel.close()
            return r
        return wrapper