from contextlib import contextmanager
from paramiko import AuthenticationException
import pysftp

@contextmanager
def open_sftp(link, login, mode='r'):
    """Context manager to read/write a file via SFTP

    :param link: Path to file
    :param login: SFTP username/password info in the format

                    {
                        'username': "Username",
                        'password': "p@ssw0rd"
                    }

    :param mode: The mode to open the file with. Either 'r' or 'w'
    :returns: File object for the specified file
    :raises:
        :ValueError: If there is an error with the input
        :AuthenticationException: If could not connect
        :IOError: If file cannot be found/opened
    """
    username = login.get('username')
    password = login.get('password')

    # Parse link into host, directory, and file
    # Not much error checking, depends on link matching:
    #   [protocol]://[host]/[path (optional)]/[file]

    # Remove the sftp protocol from the link string
    link = link.replace(u"sftp://", u"")

    link_split = link.split('/')
    if len(link_split) < 2:
        raise ValueError("Error parsing link")
    link_file = link_split.pop()
    link_path = link_split
    if not link_file:
        raise ValueError("Error parsing link: No file path")

    cinfo = {
        'host': link_split.pop(0),
        'username': username,
        'password': password
    }
    if not cinfo.get('host'):
        raise ValueError("Error parsing link: No host")
    if u":" in cinfo['host']:
        cinfo['host'], cinfo['port'] = cinfo['host'].split(u':')
        cinfo['port'] = int(cinfo['port'])

    # Connect to host. Sometimes it takes a couple tries... =(
    sftp = None
    for i in xrange(5):
        try:
            sftp = pysftp.Connection(**cinfo)
            break
        except AuthenticationException:
            warnings.warn("Could not authenticate... Will re-attempt")
            import time as t
            t.sleep(5)

    # Try one more time... if this fails, raise the AuthenticationError
    if not sftp:
        sftp = pysftp.Connection(**cinfo)

    try:
        for directory in link_path:
            try:
                sftp.chdir(directory)
            except IOError as exc:
                raise IOError("Could not navigate to file: {}".format(exc))

        try:
            file_obj = sftp.open(link_file, mode=mode)
        except IOError as exc:
            raise IOError("Could not open file: {}".format(exc))

        yield file_obj
    finally:
        sftp.close()