import json
import pathlib
import socket

from ftplib import FTP
from paramiko import Transport, SFTPClient, RSAKey
from pathlib import Path
from ssh2.session import Session
from ssh2.sftp import LIBSSH2_FXF_READ, LIBSSH2_SFTP_S_IRUSR, LIBSSH2_SFTP_S_IFREG
from typing import Union


def load_credentials(name: str):
    return json.load((Path(".credentials") / f"{name}.json").open())


class Client:
    @classmethod
    def from_creds(cls, name: str):
        args = json.load((Path(".credentials") / f"{name}.json").open())
        return cls(**args)


class FtpClient(Client):
    def __init__(self, host, username, password, folder=None):
        self.client = FTP(host, username, password)
        if folder is not None:
            self.client.cwd(folder)

    def put(self, src: Union[pathlib.Path, bytes], dst: str = None):
        if isinstance(src, pathlib.Path):
            if dst is None:
                dst = src.name
            with src.open("rb") as fh:
                self.client.storbinary(f"STOR {dst}", fh)
        else:
            self.client.storbinary(f"STOR {dst}", fh)


class SftpClient(Client):
    def __init__(self, host, port, username, password=None, key=None, folder=None):
        transport = Transport((host, port))
        if key is not None:
            pkey = RSAKey.from_private_key_file(Path.home() / ".ssh" / key)
        else:
            pkey = None
        if host.endswith("gs.com"):
            so = transport.get_security_options()
            so.key_types += ("ssh-dss",)
        transport.connect(username=username, password=password, pkey=pkey)
        self.client = SFTPClient.from_transport(transport)
        if folder is not None:
            self.client.chdir(folder)

    def download_files(self, src: str, dst: pathlib.Path, *args):
        for f in self.client.listdir(src):
            local_file = dst / f
            if not local_file.exists():
                self.client.get(f"{src}/{f}", localpath=local_file)

    def put(self, src: Union[pathlib.Path, bytes], dst: str = None):
        if isinstance(src, pathlib.Path):
            if dst is None:
                dst = src.name
            with src.open("rb") as fh:
                self.client.putfo(fh, dst)
        else:
            self.client.putfo(src, dst)


class SftpClient2(Client):
    def __init__(self, host, port, username, password):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((host, port))
        session = Session()
        session.handshake(sock)
        session.userauth_password(username, password)
        self.client = session.sftp_init()

    def download_files(self, src: str, dst: pathlib.Path, *args):
        files = []
        with self.client.opendir(src) as fh:
            for size, buf, attrs in fh.readdir():
                if attrs.permissions & LIBSSH2_SFTP_S_IFREG:
                    files.append(buf.decode())
        for f in files:
            local_file = dst / f
            if not local_file.exists():
                with self.client.open(
                    f"{src}/{f}", LIBSSH2_FXF_READ, LIBSSH2_SFTP_S_IRUSR
                ) as remote_handle, local_file.open("wb") as local_handle:
                    for size, data in remote_handle:
                        local_handle.write(data)