import os
from contextlib import contextmanager
from unittest import mock
import io

import boto3
from moto import mock_aws
import pytest
from vector_utils.connections import connection_info
from vector_utils.connections import s3


@contextmanager
def create_bucket_and_setup_files(conn_info):
    with mock_aws() as mock_aws_wrapper:
        conn = boto3.resource('s3', region_name='us-east-1')
        conn.create_bucket(Bucket='test_bucket')
        f = io.BytesIO(b'some initial binary data: \x00\x01')
        conn.meta.client.upload_fileobj(
            f, conn_info['domain_name'], '/a_dir/test.txt')
        yield mock_aws_wrapper


def test_s3_file_exists(s3_conn_info):
    """Test s3 file exists."""
    with create_bucket_and_setup_files(s3_conn_info):
        conn_obj = connection_info.ConnectionInfo(s3_conn_info)
        c = s3.S3Connection(conn_obj)

        assert c.file_exists('/a_dir/test.txt') is True
        assert c.file_exists('/a_dir/test2.txt') is False


def test_s3_connection_file_size(s3_conn_info):
    """Test s3 file size."""
    with create_bucket_and_setup_files(s3_conn_info):
        conn_obj = connection_info.ConnectionInfo(s3_conn_info)
        c = s3.S3Connection(conn_obj)

        assert c.file_size('/a_dir/test.txt') == 28


def test_s3_connection_mkdir(s3_conn_info):
    """Test s3 mkdir."""
    with create_bucket_and_setup_files(s3_conn_info):
        conn_obj = connection_info.ConnectionInfo(s3_conn_info)
        c = s3.S3Connection(conn_obj)

        assert c.mkdir('/a_dir') is None
        assert c.mkdir('/a_dir/') is None
        assert c.mkdir('/a_dir/b_dir/c_dir/d_dir/e_dir') is True
        assert c.scan_dir('/a_dir/b_dir/c_dir/d_dir') == ['e_dir']


def test_s3_connection_exception_raised(s3_conn_info):
    """Test s3 file size not found raises exception."""
    with create_bucket_and_setup_files(s3_conn_info):
        conn_obj = connection_info.ConnectionInfo(s3_conn_info)
        c = s3.S3Connection(conn_obj)

        with pytest.raises(FileNotFoundError) as excinfo:
            c.file_size('/a_dir/somefile111.txt')

        assert 'No such file' in str(excinfo)


def test_s3_connection_transfer_files_upload(local_test_files, s3_conn_info):
    """Test s3 connection class uploads files."""
    with create_bucket_and_setup_files(s3_conn_info):
        conn_obj = connection_info.ConnectionInfo(s3_conn_info)
        c = s3.S3Connection(conn_obj)

        transfer_files_list = [
            {'local': x, 'remote': '{}/{}'.format('', os.path.basename(x))}
            for x in local_test_files]

        c.transfer_files(transfer_files_list)

        assert c.scan_dir('') == [
            'a_dir', *[os.path.basename(x) for x in local_test_files]]


def test_s3_connection_transfer_files_download(tmpdir, s3_conn_info):
    """Test s3 connection class downloads files."""
    with create_bucket_and_setup_files(s3_conn_info):
        conn_obj = connection_info.ConnectionInfo(s3_conn_info)
        c = s3.S3Connection(conn_obj)

        local_file_location = '{}/{}'.format(tmpdir, 'somefile.txt')

        transfer_files_list = [{
            'local': local_file_location,
            'remote': '/a_dir/test.txt'
        }]

        c.transfer_files(transfer_files_list, transfer_mode='download')

        assert tmpdir.listdir() == [local_file_location]


@mock.patch('vector_utils.connections.s3.sleep')
def test_s3_connection_transfer_files_download_err(
        mock_sleep, tmpdir, s3_conn_info):
    """Test s3 connection class downloads files error."""
    with create_bucket_and_setup_files(s3_conn_info):
        conn_obj = connection_info.ConnectionInfo(s3_conn_info)
        c = s3.S3Connection(conn_obj)

        local_file_location = '{}/{}'.format(tmpdir, 'somefile.txt')

        transfer_files_list = [{
            'local': local_file_location,
            'remote': '/a_dir/somefile_1.txt'
        }]

        with pytest.raises(Exception) as excinfo:
            c.transfer_files(transfer_files_list, transfer_mode='download')

        assert str(excinfo.value) == 'Unable to transfer files'