import os

import posixpath
from six.moves import urllib

from mlflow.entities import FileInfo
from mlflow.store.artifact.artifact_repo import ArtifactRepository
from mlflow.exceptions import MlflowException


class SFTPArtifactRepository(ArtifactRepository):
    """Stores artifacts as files in a remote directory, via sftp."""

    def __init__(self, artifact_uri, client=None):
        self.uri = artifact_uri
        parsed = urllib.parse.urlparse(artifact_uri)
        self.config = {
            "host": parsed.hostname,
            "port": parsed.port,
            "username": parsed.username,
            "password": parsed.password,
        }
        self.path = parsed.path

        if client:
            self.sftp = client
        else:
            import pysftp
            import paramiko

            if self.config["host"] is None:
                self.config["host"] = "localhost"

            ssh_config = paramiko.SSHConfig()
            user_config_file = os.path.expanduser("~/.ssh/config")
            if os.path.exists(user_config_file):
                with open(user_config_file) as f:
                    ssh_config.parse(f)

            user_config = ssh_config.lookup(self.config["host"])

            if "hostname" in user_config:
                self.config["host"] = user_config["hostname"]

            if self.config.get("username", None) is None and "user" in user_config:
                self.config["username"] = user_config["user"]

            if self.config.get("port", None) is None:
                if "port" in user_config:
                    self.config["port"] = int(user_config["port"])
                else:
                    self.config["port"] = 22

            if "identityfile" in user_config:
                self.config["private_key"] = user_config["identityfile"][0]

            self.sftp = pysftp.Connection(**self.config)

        super(SFTPArtifactRepository, self).__init__(artifact_uri)

    def log_artifact(self, local_file, artifact_path=None):
        artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
        self.sftp.makedirs(artifact_dir)
        self.sftp.put(local_file, posixpath.join(artifact_dir, os.path.basename(local_file)))

    def log_artifacts(self, local_dir, artifact_path=None):
        artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
        self.sftp.makedirs(artifact_dir)
        self.sftp.put_r(local_dir, artifact_dir)

    def _is_directory(self, artifact_path):
        artifact_dir = self.path
        path = posixpath.join(artifact_dir, artifact_path) if artifact_path else artifact_dir
        return self.sftp.isdir(path)

    def list_artifacts(self, path=None):
        artifact_dir = self.path
        list_dir = posixpath.join(artifact_dir, path) if path else artifact_dir
        if not self.sftp.isdir(list_dir):
            return []
        artifact_files = self.sftp.listdir(list_dir)
        infos = []
        for file_name in artifact_files:
            file_path = file_name if path is None else posixpath.join(path, file_name)
            full_file_path = posixpath.join(list_dir, file_name)
            if self.sftp.isdir(full_file_path):
                infos.append(FileInfo(file_path, True, None))
            else:
                infos.append(FileInfo(file_path, False, self.sftp.stat(full_file_path).st_size))
        return infos

    def _download_file(self, remote_file_path, local_path):
        remote_full_path = posixpath.join(self.path, remote_file_path)
        self.sftp.get(remote_full_path, local_path)

    def delete_artifacts(self, artifact_path=None):
        raise MlflowException("Not implemented yet")

@@ -2,16 +2,27 @@
 mport os
+%0Aimport sys
 %0A%0Aimport
@@ -214,16 +214,533 @@
 ption%0A%0A%0A
+# Based on: https://stackoverflow.com/a/58466685%0Adef _put_r_for_windows(sftp, local_dir, remote_dir, preserve_mtime=False):%0A    for entry in os.listdir(local_dir):%0A        local_path = os.path.join(local_dir, entry)%0A        remote_path = posixpath.join(remote_dir, entry)%0A        if os.path.isdir(local_path):%0A            sftp.mkdir(remote_path)%0A            _put_r_for_windows(sftp, local_path, remote_path, preserve_mtime)%0A        else:%0A            sftp.put(local_path, remote_path, preserve_mtime=preserve_mtime)%0A%0A%0A
 class SF
@@ -2952,32 +2952,153 @@
 s(artifact_dir)%0A
+        if sys.platform == %22win32%22:%0A            _put_r_for_windows(self.sftp, local_dir, artifact_dir)%0A        else:%0A    
         self.sft

Fix bug for sftp_artifact_repo in windows (#3168)