import json
import pathlib
import socket

from ftplib import FTP
from io import BytesIO
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}", BytesIO(src))


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, confirm=True):
        if isinstance(src, pathlib.Path):
            if dst is None:
                dst = src.name
            with src.open("rb") as fh:
                self.client.putfo(fh, dst, confirm=confirm)
        else:
            self.client.putfo(BytesIO(src), dst, confirm=confirm)


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)