from unittest import TestCase

import os
import tempfile
from boto3.resources.base import ServiceResource
from botocore.client import BaseClient
from moto import mock_s3

from dblue_stores.exceptions import DblueStoresException
from dblue_stores.stores.s3 import S3Store


class TestAwsStore(TestCase):
    @mock_s3
    def test_store_client(self):
        store = S3Store(client='foo', resource='bar')
        assert store.client == 'foo'
        assert store.resource == 'bar'

        store = S3Store()
        assert isinstance(store.client, BaseClient)
        assert isinstance(store.resource, ServiceResource)

    def test_parse_s3_url(self):
        s3_url = 's3://test/this/is/valid/key.txt'
        parsed_url = S3Store.parse_s3_url(s3_url)
        assert parsed_url == ('test', 'this/is/valid/key.txt')

    def test_check_prefix_format(self):
        assert S3Store.check_prefix_format(prefix='foo', delimiter='') == 'foo'
        assert S3Store.check_prefix_format(prefix='foo', delimiter='/') == 'foo/'
        assert S3Store.check_prefix_format(prefix='foo/', delimiter='/') == 'foo/'
        assert S3Store.check_prefix_format(prefix='/foo/', delimiter='/') == '/foo/'
        assert S3Store.check_prefix_format(prefix='/foo/boo', delimiter='/') == '/foo/boo/'
        assert S3Store.check_prefix_format(prefix='', delimiter='/') == ''

    @mock_s3
    def test_check_bucket(self):
        store = S3Store()
        bucket = store.get_bucket('bucket')
        bucket.create()

        self.assertTrue(store.check_bucket('bucket'))
        self.assertFalse(store.check_bucket('not-a-bucket'))

    def test_check_bucket_raises_with_invalid_client_resources(self):
        store = S3Store(client='foo', resource='bar')

        with self.assertRaises(Exception):
            store.check_bucket('bucket')

    @mock_s3
    def test_get_bucket(self):
        store = S3Store()
        bucket = store.get_bucket('bucket')
        assert bucket is not None

    @mock_s3
    def test_list_size(self):
        store = S3Store()
        b = store.get_bucket('bucket')
        b.create()
        num_objects = 1001
        for k in range(num_objects):
            b.put_object(Key='a' + str(k), Body=b'a')

        assert len(store.list(bucket_name='bucket', delimiter='/')['keys']) == num_objects

    @mock_s3
    def test_list_prefixes(self):
        store = S3Store()
        b = store.get_bucket('bucket')
        b.create()
        b.put_object(Key='a', Body=b'a')
        b.put_object(Key='dir/b', Body=b'b')

        assert store.list_prefixes(bucket_name='bucket', prefix='non-existent/') == []
        assert store.list_prefixes(bucket_name='bucket', delimiter='/') == ['dir']
        assert store.list_keys(bucket_name='bucket', delimiter='/') == [('a', 1)]
        assert store.list_keys(bucket_name='bucket', prefix='dir/') == [('b', 1)]

    @mock_s3
    def test_ls(self):
        store = S3Store()
        b = store.get_bucket('bucket')
        b.create()
        b.put_object(Key='a', Body=b'a')
        b.put_object(Key='dir/b', Body=b'b')

        full_response = {'files': [('a', 1)], 'dirs': ['dir']}
        empty_response = {'dirs': [], 'files': []}
        dir_response = {'dirs': [], 'files': [('b', 1)]}

        assert store.ls('s3://bucket') == full_response
        assert store.ls('s3://bucket/') == full_response
        assert store.ls('s3://bucket/non-existent') == empty_response
        assert store.ls('s3://bucket/non-existent/') == empty_response
        assert store.ls('s3://bucket/dir') == dir_response
        assert store.ls('s3://bucket/dir/') == dir_response

    @mock_s3
    def test_delete(self):
        store = S3Store()
        b = store.get_bucket('bucket')
        b.create()
        b.put_object(Key='a', Body=b'a')
        b.put_object(Key='dir/b', Body=b'b')

        store.delete(key='a', bucket_name='bucket')
        assert store.list_prefixes(bucket_name='bucket', prefix='non-existent/') == []
        assert store.list_prefixes(bucket_name='bucket', delimiter='/') == ['dir']
        assert store.list_keys(bucket_name='bucket', delimiter='/') == []
        assert store.list_keys(bucket_name='bucket', prefix='dir/') == [('b', 1)]

    @mock_s3
    def test_list_prefixes_paged(self):
        store = S3Store()
        b = store.get_bucket('bucket')
        b.create()
        # Test only one page
        keys = ['x/b', 'y/b']
        dirs = ['x', 'y']
        for key in keys:
            b.put_object(Key=key, Body=b'a')

        self.assertListEqual(sorted(dirs),
                             sorted(store.list_prefixes(bucket_name='bucket',
                                                        delimiter='/',
                                                        page_size=1)))

    @mock_s3
    def test_list_keys(self):
        store = S3Store()
        b = store.get_bucket('bucket')
        b.create()
        b.put_object(Key='a', Body=b'a')
        b.put_object(Key='dir/b', Body=b'b')

        assert store.list_keys(bucket_name='bucket', prefix='non-existent/') == []
        assert store.list_keys(bucket_name='bucket') == [('a', 1), ('dir/b', 1)]
        assert store.list_keys(bucket_name='bucket', delimiter='/') == [('a', 1)]
        assert store.list_keys(bucket_name='bucket', prefix='dir/') == [('b', 1)]

    @mock_s3
    def test_list_keys_paged(self):
        store = S3Store()
        b = store.get_bucket('bucket')
        b.create()

        keys = ['x', 'y']
        for key in keys:
            b.put_object(Key=key, Body=b'a')

        self.assertListEqual(sorted([(k, 1) for k in keys]),
                             sorted(store.list_keys(bucket_name='bucket',
                                                    delimiter='/',
                                                    page_size=2)))

    @mock_s3
    def test_check_key(self):
        store = S3Store()
        b = store.get_bucket('bucket')
        b.create()
        b.put_object(Key='a', Body=b'a')

        assert store.check_key('a', 'bucket') is True
        assert store.check_key('s3://bucket//a') is True
        assert store.check_key('b', 'bucket') is False
        assert store.check_key('s3://bucket//b') is False

    @mock_s3
    def test_get_key(self):
        store = S3Store()
        b = store.get_bucket('bucket')
        b.create()
        b.put_object(Key='a', Body=b'a')

        assert store.get_key('a', 'bucket').key == 'a'
        assert store.get_key('s3://bucket/a').key == 'a'

        # No bucket
        with self.assertRaises(DblueStoresException):
            store.get_key('a', 'nobucket')

    @mock_s3
    def test_read_key(self):
        store = S3Store()
        store.client.create_bucket(Bucket='bucket')
        store.client.put_object(Bucket='bucket', Key='my_key', Body=b'M\xC3\xA9nar')

        self.assertEqual(store.read_key('my_key', 'bucket'), u'Ménar')

    @mock_s3
    def test_upload_string(self):
        store = S3Store()
        store.client.create_bucket(Bucket='bucket')

        store.upload_string(u'Ménar', 'my_key', 'bucket')
        body = store.resource.Object('bucket', 'my_key').get()['Body'].read()

        self.assertEqual(body, b'M\xC3\xA9nar')

    @mock_s3
    def test_upload_bytes(self):
        store = S3Store()
        store.client.create_bucket(Bucket='bucket')

        store.upload_bytes(b'Content', 'my_key', 'bucket')
        body = store.resource.Object('bucket', 'my_key').get()['Body'].read()

        self.assertEqual(body, b'Content')

    @mock_s3
    def test_upload_download(self):
        store = S3Store()
        store.client.create_bucket(Bucket='bucket')

        dirname = tempfile.mkdtemp()
        fpath1 = dirname + '/test1.txt'
        with open(fpath1, 'w') as f:
            f.write('data1')

        fpath2 = dirname + '/test2.txt'
        with open(fpath2, 'w') as f:
            f.write('data2')

        fpath3 = dirname + '/test3.txt'
        with open(fpath3, 'w') as f:
            f.write('data3')

        # Upload
        store.upload_file(fpath1, 'my_key1.txt', 'bucket', use_basename=False)
        assert store.check_key('my_key1.txt', 'bucket') is True

        store.upload_file(fpath2, 'my_key2.txt', 'bucket', use_basename=False)
        assert store.check_key('my_key2.txt', 'bucket') is True

        store.upload_file(fpath3, 'foo/', 'bucket', use_basename=True)
        assert store.check_key('foo/test3.txt', 'bucket') is True

        # Download
        store.download_file('my_key1.txt',
                            local_path=dirname + '/foo1.txt',
                            bucket_name='bucket',
                            use_basename=False)
        assert os.path.basename(dirname + '/foo1.txt') == 'foo1.txt'
        assert open(os.path.join(dirname + '/foo1.txt')).read() == 'data1'

        dirname2 = tempfile.mkdtemp()
        store.download_file('foo/test3.txt',
                            local_path=dirname2,
                            bucket_name='bucket',
                            use_basename=True)
        assert os.path.basename(dirname2 + '/test3.txt') == 'test3.txt'
        assert open(os.path.join(dirname2 + '/test3.txt')).read() == 'data3'

    @mock_s3
    def test_upload_download_directory(self):
        store = S3Store()
        store.client.create_bucket(Bucket='bucket')

        dirname1 = tempfile.mkdtemp()
        fpath1 = dirname1 + '/test1.txt'
        with open(fpath1, 'w') as f:
            f.write('data1')

        fpath2 = dirname1 + '/test2.txt'
        with open(fpath2, 'w') as f:
            f.write('data2')

        dirname2 = tempfile.mkdtemp(prefix=dirname1 + '/')
        fpath3 = dirname2 + '/test3.txt'
        with open(fpath3, 'w') as f:
            f.write('data3')

        rel_path1 = dirname1.split('/')[-1]
        rel_path2 = dirname2.split('/')[-1]

        # Test without using basename
        # Upload
        store.upload_dir(dirname1, 'mykey', 'bucket', use_basename=False)
        assert store.check_key('mykey/test1.txt', 'bucket') is True
        assert store.check_key('mykey/test2.txt', 'bucket') is True
        assert store.check_key('mykey/{}/test3.txt'.format(rel_path2), 'bucket') is True
        # Download
        dirname3 = tempfile.mkdtemp()
        store.download_dir('mykey', dirname3, 'bucket', use_basename=False)
        assert sorted(os.listdir(dirname3)) == sorted([rel_path2, 'test1.txt', 'test2.txt'])
        assert os.listdir('{}/{}'.format(dirname3, rel_path2)) == ['test3.txt']

        # Test with using basename
        store.upload_dir(dirname1, 'mykey', 'bucket', use_basename=True)
        assert store.check_key('mykey/{}/test1.txt'.format(rel_path1), 'bucket') is True
        assert store.check_key('mykey/{}/test2.txt'.format(rel_path1), 'bucket') is True
        assert store.check_key(
            'mykey/{}/{}/test3.txt'.format(rel_path1, rel_path2), 'bucket') is True
        # Download
        dirname3 = tempfile.mkdtemp()
        store.download_dir('mykey/{}'.format(rel_path1), dirname3, 'bucket', use_basename=True)
        assert os.listdir(dirname3) == [rel_path1]
        assert sorted(os.listdir('{}/{}'.format(dirname3, rel_path1))) == sorted(
            [rel_path2, 'test1.txt', 'test2.txt'])
        assert os.listdir('{}/{}/{}'.format(dirname3, rel_path1, rel_path2)) == ['test3.txt']
