# SPDX-License-Identifier: BSD-2-Clause
# SPDX-FileCopyrightText: 2023 KDE e.V.
# SPDX-FileContributor: Ingo Klöcker <dev@ingo-kloecker.de>

import io
import logging
import os
import stat
import time
from functools import wraps
from pathlib import Path, PurePath
from typing import Optional, Union

from paramiko import RSAKey, SSHClient
from paramiko.client import AutoAddPolicy
from paramiko.common import o777
from paramiko.ssh_exception import SSHException

log = logging.getLogger(__name__)


class SFTPError(Exception):
    pass


class ConnectionLostError(SFTPError):
    pass


class ClosingContextManager:
    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.close()


def ProgressLogger(message: str, interval: int = 1):
    lastLog = time.time()

    def printProgress(current: int, total: int):
        nonlocal lastLog
        now = time.time()
        if current == total or now > lastLog + interval:
            lastLog = now
            log.debug(message, current, total)

    return printProgress


def toSftpPath(path: Union[str, PurePath]):
    if isinstance(path, PurePath):
        return str(path.as_posix())
    return path


def detectLostConnection(f):
    @wraps(f)
    def wrapper(self, *args, **kwargs):
        try:
            return f(self, *args, **kwargs)
        except OSError:
            # probably a socket error; check if the connection was closed
            if self.closed:
                raise ConnectionLostError("Lost connection to SFTP server.")
            else:
                raise

    return wrapper


class SFTPClient(ClosingContextManager):
    def __init__(self):
        super().__init__()

        self._connectInfo = None
        self._sshClient = None
        self._sftpClient = None
        self.basePath = None

    @property
    def sftpClient(self):
        if self._sftpClient is None:
            self._sftpClient = self._sshClient.open_sftp()
            if self.basePath is not None:
                try:
                    path = toSftpPath(self.basePath)
                    self._sftpClient.chdir(path)
                    path = None
                except FileNotFoundError:
                    log.debug(f"Folder '{path}' does not exist. Trying to create it ...")
                if path is not None:
                    try:
                        self._sftpClient.mkdir(path)
                        self._sftpClient.chdir(path)
                    except OSError as e:
                        raise SFTPError(
                            f"Failed to create folder '{path}' on the SFTP server ({e!r}). Create it manually."
                        )
        return self._sftpClient

    def close(self):
        if self._sftpClient is not None:
            self._sftpClient.close()
            self._sftpClient = None
        if self._sshClient is not None:
            self._sshClient.close()
            self._sshClient = None

    @property
    def closed(self):
        return self._sftpClient is None or self._sftpClient.get_channel()._closed

    def connect(
        self,
        hostname: str,
        basePath: Union[str, PurePath],
        knownHostsFile: Optional[Union[str, Path]] = None,
        rsaKeyFile: Optional[Union[str, Path]] = None,
        **kwargs,
    ):
        self._connectInfo = {
            "args": (hostname, basePath),
            "kwargs": {**kwargs, "knownHostsFile": knownHostsFile, "rsaKeyFile": rsaKeyFile},
        }

        self.basePath = basePath

        connectArgs = dict(**kwargs)
        if rsaKeyFile is not None:
            if isinstance(rsaKeyFile, Path):
                rsaKeyFile = str(rsaKeyFile)
            # explicitly load the RSA key file because using SSHClient.connect with key_file fails
            # (https://github.com/paramiko/paramiko/issues/1839)
            if isinstance(rsaKeyFile, str) and rsaKeyFile.startswith("env:"):
                envVarValue = os.environ[rsaKeyFile[4:]]
                # guess whether it's a file path or an actual RSA key
                if envVarValue.startswith("-----"):
                    # looks like a private (RSA) key
                    privateKey = io.StringIO(envVarValue)
                    connectArgs["pkey"] = RSAKey.from_private_key(privateKey)
                else:
                    connectArgs["pkey"] = RSAKey.from_private_key_file(envVarValue)
            else:
                connectArgs["pkey"] = RSAKey.from_private_key_file(rsaKeyFile)
            connectArgs["allow_agent"] = False  # use agent only if we do not have a key
        self._sshClient = SSHClient()
        if isinstance(knownHostsFile, Path):
            knownHostsFile = str(knownHostsFile)
        self._sshClient.load_system_host_keys(knownHostsFile)
        if knownHostsFile is None:
            # Set "nosec B507" as this is intentional behaviour
            self._sshClient.set_missing_host_key_policy(AutoAddPolicy())  # nosec B507
        self._sshClient.connect(hostname, **connectArgs)
        self.sftpClient

    def reconnect(self, maxTries: int = -1, timeout: int = 30 * 60):
        if self._connectInfo is None:
            # connect() was never called
            log.warning("Connect info is None. Ignoring reconnect request.")
            return

        # ensure that the previous connection is closed properly
        self.close()

        tries = 0
        endTime = time.time() + timeout
        maxCoolingOffPeriod = 60
        coolingOffPeriod = 1.0
        while True:
            tries += 1
            e: Optional[Exception] = None
            try:
                self.connect(*self._connectInfo["args"], **self._connectInfo["kwargs"])
                return
            except SSHException as ex:
                e = ex
            except OSError as ex:
                e = ex
            if maxTries > 0 and tries >= maxTries:
                raise SFTPError(f"Failed to reconnect to the SFTP server ({e!r}). Giving up after {maxTries} retries.")
            now = time.time()
            if now > endTime:
                raise SFTPError(f"Failed to reconnect to the SFTP server ({e!r}). Giving up after {timeout} seconds.")
            coolingOffPeriod = min(coolingOffPeriod, endTime - now)
            log.debug(f"Failed to reconnect to the SFTP server ({e!r}). Retrying in {coolingOffPeriod} s ...")
            time.sleep(coolingOffPeriod)
            coolingOffPeriod = min(coolingOffPeriod * 2, maxCoolingOffPeriod)

    def download(self, remotePath: Union[str, PurePath], localPath: Union[str, Path]):
        if isinstance(localPath, str):
            localPath = Path(localPath)
        log.info(f"Downloading file {remotePath!s} to {localPath!s} ...")
        try:
            self.sftpClient.get(
                toSftpPath(remotePath), toSftpPath(localPath), ProgressLogger("Downloaded %s bytes of %s bytes")
            )
        except Exception as e:
            # remove a potentially existing partially downloaded file
            localPath.unlink(missing_ok=True)
            raise SFTPError(f"Failed to download file '{remotePath!s}' from the SFTP server ({e!r})")

    def downloadBytes(self, remotePath: Union[str, PurePath]):
        log.info(f"Downloading file {remotePath!s} ...")
        try:
            with io.BytesIO() as dataStream:
                self.sftpClient.getfo(
                    toSftpPath(remotePath), dataStream, callback=ProgressLogger("Downloaded %s bytes of %s bytes")
                )
                return dataStream.getvalue()
        except Exception as e:
            raise SFTPError(f"Failed to download file '{remotePath!s}' from the SFTP server ({e!r})")

    def upload(self, localPath: Union[str, PurePath], remotePath: Union[str, PurePath]):
        log.info(f"Uploading file {localPath!s} to {remotePath!s} ...")

        maxTries = 3
        tries = 0
        while True:
            exception = None
            try:
                self.sftpClient.put(
                    toSftpPath(localPath), toSftpPath(remotePath), ProgressLogger("Uploaded %s bytes of %s bytes")
                )
                return
            except EOFError as e:
                tries += 1
                exception = e
            except Exception as e:
                raise SFTPError(f"Failed to upload file '{localPath!s}' to the SFTP server ({e!r})")

            if tries < maxTries:
                log.error("Error: Uploading file '%s' failed with %r. Retrying ...", localPath, exception)
                self.reconnect()
                try:
                    self.sftpClient.remove(toSftpPath(remotePath))
                except FileNotFoundError:
                    pass
            else:
                raise SFTPError(f"Failed to upload file '{localPath!s}' to the SFTP server ({exception!r})")

    def uploadBytes(self, data: bytes, remotePath: Union[str, PurePath]):
        dataSize = len(data)
        log.info(f"Uploading {dataSize} bytes to {remotePath!s} ...")

        maxTries = 3
        tries = 0
        while True:
            dataStream = io.BytesIO(data)
            exception = None
            try:
                self.sftpClient.putfo(
                    dataStream,
                    toSftpPath(remotePath),
                    file_size=dataSize,
                    callback=ProgressLogger("Uploaded %s bytes of %s bytes"),
                )
                return
            except EOFError as e:
                tries += 1
                exception = e
            except Exception as e:
                raise SFTPError(f"Failed to upload data to the SFTP server ({e!r})")

            if tries < maxTries:
                log.error(f"Error: Uploading data failed with {exception!r}. Retrying ...")
                self.reconnect()
                try:
                    self.sftpClient.remove(toSftpPath(remotePath))
                except FileNotFoundError:
                    pass
            else:
                raise SFTPError(f"Failed to upload data to the SFTP server ({exception!r})")

    @detectLostConnection
    def listdir(self, path: Union[str, PurePath] = "."):
        return self.sftpClient.listdir(toSftpPath(path))

    @detectLostConnection
    def listdir_attr(self, path: Union[str, PurePath] = "."):
        return self.sftpClient.listdir_attr(toSftpPath(path))

    @detectLostConnection
    def listdir_iter(self, path: Union[str, PurePath] = ".", read_aheads=50):
        return self.sftpClient.listdir_iter(toSftpPath(path), read_aheads=read_aheads)

    @detectLostConnection
    def lstat(self, path: Union[str, PurePath]):
        return self.sftpClient.lstat(toSftpPath(path))

    @detectLostConnection
    def mkdir(self, path: Union[str, PurePath], mode=o777):
        return self.sftpClient.mkdir(toSftpPath(path), mode=mode)

    @detectLostConnection
    def posix_rename(self, oldpath: Union[str, PurePath], newpath: Union[str, PurePath]):
        return self.sftpClient.posix_rename(toSftpPath(oldpath), toSftpPath(newpath))

    @detectLostConnection
    def remove(self, path: Union[str, PurePath]):
        return self.sftpClient.remove(toSftpPath(path))

    @detectLostConnection
    def rmdir(self, path: Union[str, PurePath]):
        return self.sftpClient.rmdir(toSftpPath(path))

    def recursiveRemoveDirectory(self, path: Union[str, PurePath]):
        path = PurePath(path)
        entriesAttributes = self.listdir_attr(path)
        for entryAttributes in entriesAttributes:
            entryPath = path / entryAttributes.filename
            if stat.S_ISDIR(entryAttributes.st_mode):
                self.recursiveRemoveDirectory(entryPath)
            else:
                self.remove(entryPath)
        self.rmdir(path)

    @detectLostConnection
    def stat(self, path: Union[str, PurePath]):
        return self.sftpClient.stat(toSftpPath(path))