import csv
import gzip
import tempfile
import os
import io
import logging
import unittest
import unittest.mock

from s3_uploader.copier import universal


class MockedCursor:
    def __init__(self, headers: list, data: list):
        self.description = [(header,) for header in headers]
        self.iterator = iter(data)

    def fetchone(self):
        try:
            return next(self.iterator)
        except StopIteration:
            return None

    def fetchmany(self, *args, **kwargs):
        try:
            record = next(self.iterator)
            return [record, ]
        except StopIteration:
            return None


class UniversalCopierTestCase(unittest.TestCase):

    @unittest.mock.patch('smart_open.open')
    def test_raise_exception_in_thread(self, mocked_open):
        """Chech that exception is propagated"""
        mocked_open.side_effect = ConnectionError
        with tempfile.TemporaryDirectory() as tmp_dir:
            self.assertRaises(
                ConnectionError,
                universal.copy,
                logging.getLogger(),
                io.BytesIO(b'some data'),
                workers=[
                    universal.UploadWorker(
                        logger=unittest.mock.Mock(),
                        path=os.path.join(tmp_dir, 'error.txt'),
                        ignore_ext=True,
                    )
                ],
                chunk_size=1024,
            )

    def test_stream_to_fs(self):
        """Save STREAM to TXT"""
        with tempfile.TemporaryDirectory() as tmp_dir:
            data = b'A lot of bytes here.'
            path = os.path.join(tmp_dir, 'test_report.txt')

            total_read = universal.copy(
                logging.getLogger(),
                io.BytesIO(data),
                workers=[
                    universal.UploadWorker(
                        logger=unittest.mock.Mock(),
                        path=path,
                        ignore_ext=True,
                    )
                ],
                chunk_size=1024,
            )
            self.assertEqual(total_read, len(data))
            self.assertEqual(open(path, 'rb').read(), data)

    def test_stream_to_fs_gz(self):
        """Save STREAM to GZ (archive on the fly)"""
        with tempfile.TemporaryDirectory() as tmp_dir:
            data = b'A lot of bytes here.'
            path = os.path.join(tmp_dir, 'test_report.txt.gz')

            total_read = universal.copy(
                logging.getLogger(),
                io.BytesIO(data),
                workers=[
                    universal.UploadWorker(
                        logger=unittest.mock.Mock(),
                        path=path,
                        ignore_ext=False,
                    )
                ],
                chunk_size=1024,
            )
            self.assertEqual(total_read, len(data))
            self.assertEqual(gzip.open(path, 'rb').read(), data)

    def test_stream_to_multiple(self):
        """Save stream both to TXT and GZ files."""
        with tempfile.TemporaryDirectory() as tmp_dir:
            data = b'A lot of bytes here.'
            path_txt = os.path.join(tmp_dir, 'test_report.txt')
            path_gz = os.path.join(tmp_dir, 'test_report.txt.gz')

            total_read = universal.copy(
                logging.getLogger(),
                io.BytesIO(data),
                workers=[
                    universal.UploadWorker(
                        logger=unittest.mock.Mock(),
                        path=path_txt,
                        ignore_ext=True,
                    ),
                    universal.UploadWorker(
                        logger=unittest.mock.Mock(),
                        path=path_gz,
                        ignore_ext=False,
                    )
                ],
                chunk_size=1024,
            )
            self.assertEqual(total_read, len(data))
            self.assertEqual(gzip.open(path_gz, 'rb').read(), data)
            self.assertEqual(open(path_txt, 'rb').read(), data)

    def test_cursor_to_fs_gz(self):
        """Save CSV based on cursor data to GZ"""
        with tempfile.TemporaryDirectory() as tmp_dir:
            cursor = MockedCursor(
                ['Name', 'Age'], [
                    ['John', 27],
                    ['Sara', 23],
                ])

            path = os.path.join(tmp_dir, 'report.csv.gz')

            universal.copy(
                logging.getLogger(),
                universal.CursorCSVStream(cursor),
                workers=[
                    universal.UploadWorker(
                        logger=unittest.mock.Mock(),
                        path=path,
                        ignore_ext=False,
                    )
                ],
                chunk_size=1024,
            )

            with gzip.open(path, 'rt') as fout:
                reader = csv.reader(fout)
                self.assertEqual(next(reader), ['Name', 'Age'])
                self.assertEqual(next(reader), ['John', '27'])
                self.assertEqual(next(reader), ['Sara', '23'])
                self.assertRaises(StopIteration, next, reader)


class CSVCusrosStreamTestCase(unittest.TestCase):

    def test_read_method_all(self):
        cursor = MockedCursor(['One', 'Two'], [[1, 2]])
        stream = universal.CursorCSVStream(cursor)
        self.assertEqual(stream.read(), b'One,Two\r\n1,2\r\n')

    def test_read_method_big_chunk_size(self):
        cursor = MockedCursor(['One', 'Two'], [[1, 2], [3, 4]])
        stream = universal.CursorCSVStream(cursor)
        self.assertEqual(stream.read(1024), b'One,Two\r\n1,2\r\n3,4\r\n')

    def test_read_method_N_size(self):
        cursor = MockedCursor(['Name', 'Age'], [['John', 27], ['Sara', 23]])
        stream = universal.CursorCSVStream(cursor)
        self.assertEqual(stream.read(5), b'Name,')
        self.assertEqual(stream.read(5), b'Age\r\n')