from ftplib import FTP
from ftplib import error_perm
import os
import shutil

from boto.s3.connection import Bucket
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from garcon import task  # noqa
import paramiko


def file_exist_on_sftp(path, **sftp_creds):
    """Check if file exists on sftp.

    Args:
        path (str): absolute remote path on sftp to the files.
        sftp_creds (dict): dictionary contains sftp credentials.

    Returns:
        bool: True if file exists, False otherwise
    """

    try:
        transport = paramiko.Transport(
            (sftp_creds['host'], sftp_creds['port']))
        transport.connect(
            username=sftp_creds['username'], password=sftp_creds['password'])
        sftp = paramiko.SFTPClient.from_transport(transport)
        sftp.stat(path)
        return True
    except FileNotFoundError:
        return False


def _extract_bucket_path(url):
    """Extract the bucket name and path from the provided S3 url

    Args:
        url (str): S3 URL like 's3://bucket_name/folder1/folder2/'

    Return:
        tuple: bucket name and bucket path
    """

    bucket_path = ''
    bucket_name = ''

    if url.startswith('s3://'):
        split_path = url.split('/', 3)
        bucket_name = split_path[2]

        if len(split_path) > 3:
            bucket_path = split_path[3]

    if not bucket_name:
        raise Exception('The S3 url \'{url}\' is not valid.'.format(url=url))

    return (bucket_name, bucket_path)


def _upload_from_local_to_ftp(
        sftp_creds, file_name, remote_ftp_dir, local_dir, pkey=None):
    """Upload files from local file system to FTP

    Args:
        sftp_creds (dict): credentials for connecting to ftp server
        file_name (str): name of the file to be uploaded
        remote_ftp_dir (str): destination directory on FTP
        local_dir (str): local directory where the file is located
        pkey (str): path to ssh private key
    """

    host = sftp_creds['host']
    port = int(sftp_creds['port'])
    username = sftp_creds['username']
    remote_ftp_dir = remote_ftp_dir.rstrip('/')

    if pkey or port != 21:
        transport = paramiko.Transport((host, port))
        if pkey:
            transport.connect(
                username=username,
                pkey=paramiko.RSAKey.from_private_key_file(pkey))
        else:
            transport.connect(
                username=username,
                password=sftp_creds['password'])
        sftp = paramiko.SFTPClient.from_transport(transport)
        dir_list = remote_ftp_dir.strip('/').split('/')
        full_path = ''

        for current_dir in dir_list:
            try:
                full_path = '{previous}/{current}'.format(
                    previous=full_path, current=current_dir)
                sftp.chdir(full_path)
            except (IOError, FileExistsError):
                sftp.mkdir(full_path)
                sftp.chdir(full_path)
        sftp.put(
            '{}/{}'.format(local_dir, file_name),
            '/{}/{}'.format(remote_ftp_dir, file_name))
        sftp.close()
        transport.close()
    else:
        ftp = FTP(host)
        ftp.login(username, sftp_creds['password'])
        try:
            ftp.mkd(remote_ftp_dir)
        except Exception:
            pass
        ftp.cwd(remote_ftp_dir)

        with open('{}/{}'.format(local_dir, file_name), 'rb') as fp:
            ftp.storbinary(
                'STOR {}'.format(file_name), fp, 1024)
        ftp.close()


def _download_from_ftp(ftp_creds, file_name, remote_ftp_dir, local_dir):
    """Download file from SFTP to local file system.

    Args:
        ftp_creds (dict): credentials for connecting to ftp server.
        file_name (str): name of the file to be uploaded.
        remote_ftp_dir (str): source directory on FTP.
        local_dir (str): local destination directory where the file
            will be stored.
    """

    remote_ftp_dir = remote_ftp_dir.rstrip('/')

    ftp = FTP(ftp_creds['host'])
    ftp.login(ftp_creds['username'], ftp_creds['password'])
    ftp.cwd(remote_ftp_dir)

    with open('{}/{}'.format(local_dir, file_name), 'wb') as fp:
        ftp.retrbinary('RETR {}'.format(file_name), fp.write, 1024)
    ftp.close()


def _download_from_sftp(
        sftp_creds, file_name, remote_ftp_dir, local_dir, pkey=None):
    """Download file from SFTP to local file system.

    Args:
        sftp_creds (dict): credentials for connecting to SFTP server.
        file_name (str): name of the file to be uploaded.
        remote_ftp_dir (str): source directory on SFTP.
        local_dir (str): local destination directory where the file
            will be stored.
        pkey (str): path to ssh private key.
    """

    host = sftp_creds['host']
    port = sftp_creds['port']
    username = sftp_creds['username']
    password = sftp_creds.get('password', None)
    remote_ftp_dir = remote_ftp_dir.rstrip('/')

    if pkey:
        pkey = paramiko.RSAKey.from_private_key_file(pkey)

    transport = paramiko.Transport((host, port))
    transport.connect(username=username, password=password, pkey=pkey)

    sftp = paramiko.SFTPClient.from_transport(transport)

    if remote_ftp_dir:  # if not empty string
        # try change working dir of the current session
        try:
            # sometimes we can and should chdir without slash
            sftp.chdir(remote_ftp_dir)
        except IOError:
            # sometimes we have to add a slash, e.g. '/GooglePlay_daily'
            # due to FTP server settings we can't chdir to 'GooglePlay_daily'
            sftp.chdir('/{}'.format(remote_ftp_dir))
        # after workdir was changed, we can't use full path
        remote_file_path = file_name
    else:
        # support current behaviour (before this commit), just in case
        remote_file_path = '/{}'.format(file_name)

    local_file_path = '{}/{}'.format(local_dir, file_name)

    # Avoid using paramiko's built in paramiko.SFTPClient.get() bc
    # Windows ftp server chokes on its custom prefetching
    # (adjusting the SFTPFile.MAX_REQUEST_SIZE also did not work)
    # http://stackoverflow.com/questions/12486623/paramiko-fails-to-download-large-files-1gb/14210233#14210233
    remote_file_size = sftp.stat(remote_file_path).st_size
    with open(local_file_path, 'wb') as local_ftp_file, sftp.open(
            remote_file_path) as remote_ftp_file:
        shutil.copyfileobj(remote_ftp_file, local_ftp_file)
    local_file_size = os.stat(local_file_path).st_size
    if local_file_size != remote_file_size:
        raise IOError(
            'size mismatch in sftp download {} != {}'.format(
                local_file_size, remote_file_size))
    sftp.close()
    transport.close()


@task.decorate(timeout=7200)
def copy_from_s3_to_ftp(
        activity, sftp_creds, file_names_list, remote_s3_dir_path,
        remote_ftp_dir_path, pkey):
    """Copy files from S3 to FTP

    Args:
        activity (ActivityWorker): The swf activity worker.
        sftp_creds (dict): sftp connection credentials dictionary
        file_names_list (list): list of file names
        remote_s3_dir_path (str): s3 path to the folder which contains the
        source file
        remote_ftp_dir_path (str): path to the destination sftp location
        pkey (Optional[str]): path to private key. If present, use key,
        otherwise use username and password
    """

    remote_s3_dir_path = remote_s3_dir_path.rstrip('/')
    bucket_name, bucket_path = _extract_bucket_path(remote_s3_dir_path)
    bucket = Bucket(S3Connection(), bucket_name)

    for file_name in file_names_list:
        local_dir = './'
        remote_dir = '{remote_ftp_path}/'.format(
            remote_ftp_path=remote_ftp_dir_path)

        key = bucket.get_key('{bucket_path}/{file_name}'.format(
            bucket_path=bucket_path, file_name=file_name))
        key.get_contents_to_filename('{}/{}'.format(local_dir, file_name))

        _upload_from_local_to_ftp(
            sftp_creds, file_name, remote_dir, local_dir, pkey)

        # Remove file from local dir when done uploading
        os.remove('{}/{}'.format(local_dir, file_name))


@task.decorate(timeout=7200)
def copy_from_ftp_to_s3(
        activity, sftp_creds, file_names_list, remote_s3_dir_path,
        remote_ftp_dir_path, pkey):
    """Copy files from FTP to S3.

    Args:
        activity (ActivityWorker): The swf activity worker.
        sftp_creds (dict): sftp connection credentials dictionary
        file_names_list (list): list of file names.
        remote_s3_dir_path (str): s3 path to the folder which contains the
            source file.
        remote_ftp_dir_path (str): path to the destination sftp location
        pkey (Optional[str]): path to private key.

    Return:
        list: List of dicts containing:
            - file: filename
            - status: true/false successful
            - key: resulting key (on success)
            - exception: exception (on exception)
    """

    remote_s3_dir_path = remote_s3_dir_path.rstrip('/')
    bucket_name, bucket_path = _extract_bucket_path(remote_s3_dir_path)
    bucket = Bucket(S3Connection(), bucket_name)
    response = dict(files=[])

    for file_name in file_names_list:
        resp = {'file': file_name, 'status': True}
        local_dir = './'
        remote_dir = '{remote_ftp_path}/'.format(
            remote_ftp_path=remote_ftp_dir_path)

        try:
            if pkey or sftp_creds['port'] != 21:
                _download_from_sftp(
                    sftp_creds, file_name, remote_dir, local_dir, pkey)
            else:
                _download_from_ftp(
                    sftp_creds, file_name, remote_dir, local_dir)
            key = Key(bucket, '{bucket_path}/{file_name}'.format(
                bucket_path=bucket_path, file_name=file_name))
            key.set_contents_from_filename('{}/{}'.format(
                local_dir, file_name))
            resp['file_size'] = key.size
            # Remove file from local dir when done uploading
            os.remove('{}/{}'.format(local_dir, file_name))
        except (error_perm, FileNotFoundError) as e:
            resp['status'] = False
            resp['exception'] = e

        response['files'].append(resp)

    return response


@task.decorate(timeout=7200)
def copy_file_from_ftp_to_s3(
        activity, ftp_creds, ftp_path, ftp_file_name, s3_path,
        s3_file_name='', pkey=None):
    """Copy files from FTP to S3.

    Args:
        activity (ActivityWorker): The swf activity worker.
        ftp_creds (dict): FTP connection credentials dictionary.
        ftp_path (str): Path to source FTP directory.
        ftp_file_name (str): Source file name.
        s3_path (str): Target path on S3.
        s3_file_name (Optional[str]): Target file name. If value is empty
        string the target name will be the same with source file name.
        pkey (Optional[str]): path to private key.

    Return:
        dict:
            - file: filename
            - status: true/false successful
            - file_size: file size in bytes
            - exception: exception (on exception)
    """

    # set correct target file name
    file_name = s3_file_name if s3_file_name else ftp_file_name
    response = {'file': file_name, 'status': True}
    local_dir = '.'

    try:
        if pkey or ftp_creds['port'] != 21:
            _download_from_sftp(
                ftp_creds, ftp_file_name, ftp_path, local_dir, pkey)
        else:
            _download_from_ftp(
                ftp_creds, ftp_file_name, ftp_path, local_dir)

        bucket_name, bucket_path = _extract_bucket_path(s3_path)
        bucket = Bucket(S3Connection(), bucket_name)
        key = Key(bucket, '{bucket_path}/{file_name}'.format(
            bucket_path=bucket_path, file_name=file_name))
        key.set_contents_from_filename('{}/{}'.format(
            local_dir, ftp_file_name))
        response['file_size'] = key.size
        # Remove file from local dir when done uploading
        os.remove('{}/{}'.format(local_dir, ftp_file_name))
    except (error_perm, FileNotFoundError) as e:
        response['status'] = False
        response['exception'] = e

    return response