import codecs
import io
import unittest
import unittest.mock
import zipfile

import boto3
from botocore.exceptions import ClientError
from moto import mock_s3

from s3_uploader.command import get_copier
from s3_uploader.copier.stream_to_s3 import Copier as StreamToS3Copier
from s3_uploader.copier.stream_to_s3_decompress import Copier as StreamToS3DecompressCopier


class StreamToS3CopierTest(unittest.TestCase):

    @mock_s3
    def test_copier_from_stream_to_s3(self):
        boto3.setup_default_session()
        conn = boto3.resource('s3', region_name='us-east-1')
        conn.create_bucket(Bucket='dev-sme-stream-transfer')

        expected = 'Some bytes data'

        copier = get_copier(
            unittest.mock.Mock(),
            io.BytesIO(bytes(expected, 'utf-8')),
            's3://dev-sme-stream-transfer/test_file.gz',
            None
        )
        copier.copy()
        actual = conn.Object('dev-sme-stream-transfer', 'test_file.gz').get()['Body'].read().decode("utf-8")
        self.assertEqual(actual, expected)

    def test_get_copier_from_stream_to_s3(self):
        copier = get_copier(
            unittest.mock.Mock(),
            io.BytesIO(bytes("Some bytes data", 'utf-8')),
            's3://dev-sme-stream-transfer/test_file.gz',
            None
        )

        self.assertIsInstance(copier, StreamToS3Copier)

    @mock_s3
    def test_copier_from_stream_to_s3_decompressed_gz(self):
        boto3.setup_default_session()
        conn = boto3.resource('s3', region_name='us-east-1')
        conn.create_bucket(Bucket='dev-sme-stream-transfer')

        expected = 'Some bytes data'
        data = bytes(expected, 'utf-8')
        expected_compressed = codecs.encode(data, "zlib")

        copier = get_copier(
            unittest.mock.Mock(),
            io.BytesIO(expected_compressed),
            's3://dev-sme-stream-transfer/test_file.gz',
            's3://dev-sme-stream-transfer/test_file.txt'
        )

        total = copier.copy()
        self.assertEqual(total, len(expected_compressed))

        actual_compressed = conn.Object('dev-sme-stream-transfer', 'test_file.gz').get()['Body'].read()
        self.assertEqual(actual_compressed, expected_compressed)

        actual_decompressed = conn.Object('dev-sme-stream-transfer', 'test_file.txt').get()['Body'].read().decode("utf-8")
        self.assertEqual(actual_decompressed, expected)

    @mock_s3
    def test_copier_from_stream_to_s3_decompressed_only_gz(self):
        boto3.setup_default_session()
        conn = boto3.resource('s3', region_name='us-east-1')
        conn.create_bucket(Bucket='dev-sme-stream-transfer')

        expected = 'Some bytes data'
        data = bytes(expected, 'utf-8')
        expected_compressed = codecs.encode(data, "zlib")

        copier = StreamToS3DecompressCopier(
            unittest.mock.Mock(),
            io.BytesIO(expected_compressed),
            {"bucket": "dev-sme-stream-transfer", "key": "test_file.gz"},
            {"bucket": "dev-sme-stream-transfer", "key": "test_file.txt"},
            decompress_only=True
        )
        total = copier.copy()
        self.assertEqual(total, len(expected_compressed))

        self.assertRaises(ClientError, conn.Object('dev-sme-stream-transfer', 'test_file.gz').get)

        actual_decompressed = conn.Object('dev-sme-stream-transfer', 'test_file.txt').get()['Body'].read().decode("utf-8")
        self.assertEqual(actual_decompressed, expected)

    @mock_s3
    def test_copier_from_stream_to_s3_decompressed_zip(self):
        boto3.setup_default_session()
        conn = boto3.resource('s3', region_name='us-east-1')
        conn.create_bucket(Bucket='dev-sme-stream-transfer')

        mem_zip = io.BytesIO()
        with zipfile.ZipFile(mem_zip, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
            zf.writestr('test_file.txt', 'test zip file data')

        expected_compressed = mem_zip.getvalue()
        copier = get_copier(
            unittest.mock.Mock(),
            io.BytesIO(expected_compressed),
            's3://dev-sme-stream-transfer/test_file.zip',
            's3://dev-sme-stream-transfer/test_file.txt'
        )

        copier.copy()

        actual_compressed = conn.Object('dev-sme-stream-transfer', 'test_file.zip').get()['Body'].read()
        self.assertEqual(actual_compressed, expected_compressed)

        actual_decompressed = conn.Object('dev-sme-stream-transfer', 'test_file.txt').get()['Body'].read().decode("utf-8")
        self.assertEqual(actual_decompressed, "test zip file data")

    @mock_s3
    def test_copier_from_stream_to_s3_decompressed_only_zip(self):
        boto3.setup_default_session()
        conn = boto3.resource('s3', region_name='us-east-1')
        conn.create_bucket(Bucket='dev-sme-stream-transfer')

        mem_zip = io.BytesIO()
        with zipfile.ZipFile(mem_zip, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
            zf.writestr('test_file.txt', 'test zip file data1')

        expected_compressed = mem_zip.getvalue()
        copier = StreamToS3DecompressCopier(
            unittest.mock.Mock(),
            io.BytesIO(expected_compressed),
            {"bucket": "dev-sme-stream-transfer", "key": "test_file.zip"},
            {"bucket": "dev-sme-stream-transfer", "key": "test_file.txt"},
            decompress_only=True
        )

        copier.copy()

        self.assertRaises(ClientError, conn.Object('dev-sme-stream-transfer', 'test_file.zip').get)

        actual_decompressed = conn.Object('dev-sme-stream-transfer', 'test_file.txt').get()['Body'].read().decode("utf-8")
        self.assertEqual(actual_decompressed, "test zip file data1")

    @mock_s3
    def test_get_copier_from_stream_to_s3_decompressed(self):
        copier = get_copier(
            unittest.mock.Mock(),
            io.BytesIO(b'Some bytes data'),
            's3://dev-sme-stream-transfer/test_file.gz',
            's3://dev-sme-stream-transfer/test_file.csv'
        )

        self.assertIsInstance(copier, StreamToS3DecompressCopier)