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 not key_path else None
        self.port = port
        self.nb_bytes = 1024
        self.keys = []
        self.transport = None
        key_type = key_type.lower()

        if key_path:
            key_file = open(path.expanduser(key_path), 'r')
            key = _KEY_TYPES[key_type].from_private_key(key_file, key_password)
            self.keys.append(key)
        elif ssh_password is None:
            self.keys = paramiko.Agent().get_keys()
            try:
                key_file = open(path.expanduser(f"~/.ssh/id_{key_type}"), 'r')
                key = _KEY_TYPES[key_type].from_private_key(
                    key_file, key_password)
            except Exception:
                pass
            else:
                index = len(self.keys) if key_password is None else 0
                self.keys.insert(index, key)
            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,
        capture=False,
        shell=True,
        combine_stderr=False,
    ):
        exit_code, output = 0, ""
        channel = self.transport.open_session()
        channel.settimeout(2)
        channel.set_combine_stderr(combine_stderr)
        if shell: channel.get_pty()
        channel.exec_command(command)

        if not display and not capture:
            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 raw_data: break
                data = raw_data.decode("utf-8")
                if display: print(data, end='')
                if capture: output += data
                if stop_event.is_set(): break

        channel.close()

        if channel.exit_status_ready():
            exit_code = channel.recv_exit_status()

        return (exit_code, output.splitlines())

    def _run_until_exit(
        self,
        command,
        timeout,
        display=True,
        capture=False,
        shell=True,
        combine_stderr=False,
    ):
        exit_code, output = 0, ""
        channel = self.transport.open_session()
        channel.settimeout(timeout)
        channel.set_combine_stderr(combine_stderr)
        if shell: channel.get_pty()
        channel.exec_command(command)

        try:
            if not display and not capture:
                return (channel.recv_exit_status(), output.splitlines())
            else:
                while True:
                    raw_data = channel.recv(self.nb_bytes)
                    if not raw_data: break
                    data = raw_data.decode("utf-8")
                    if display: print(data, end='')
                    if capture: output += data
        except socket.timeout:
            logging.warning(f"Timeout after {timeout}s")
            exit_code = 1
        except KeyboardInterrupt:
            logging.info("KeyboardInterrupt")
            exit_code = 0
        else:
            exit_code = channel.recv_exit_status()
        finally:
            channel.close()
            return (exit_code, output.splitlines())

    def run(
        self,
        command,
        display=False,
        capture=False,
        shell=True,
        combine_stderr=False,
        timeout=None,
        stop_event=None,
    ):
        if stop_event is not None:
            return self._run_until_event(
                command,
                stop_event,
                display=display,
                shell=shell,
                combine_stderr=combine_stderr,
                capture=capture,
            )
        else:
            return self._run_until_exit(
                command,
                timeout,
                display=display,
                shell=shell,
                combine_stderr=combine_stderr,
                capture=capture,
            )

    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
            sftp_channel = SFTPController.from_transport(self.transport)
            r = getattr(sftp_channel, target)(*args, **kwargs)
            sftp_channel.close()
            return r

        return wrapper