import pytest
import pytz
import uuid

from botocore.exceptions import ClientError
from datetime import datetime
from psi_kinesis.kinesis_reader import PROVISIONED_THROUGHPUT_EXCEPTION
from psi_kinesis.kinesis_writer import KinesisWriter, MAX_INVALID_RECORDS, MAX_PUT_RECORDS_BATCH
from psi_utils.unittest_mock import MockApi
from typing import Tuple


class TestKinesisWriter_UnitTests:

    @pytest.mark.unit
    def test_put_records(self):
        # Set expected values
        expected_records = [{'Data': b'Some data', 'PartitionKey': 'partitionkey'}]
        expected_stream_name = 'kinesis-stream-name'
        expected_sequence_number = '1234567890123456789'
        expected_shard_id = 'shard-000001'
        expected_resp = [{
            'SequenceNumber': expected_sequence_number,
            'ShardId': expected_shard_id
        }]
        expected_put_records_batch_resp = {'Records': expected_resp}

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'put_records',
            Records=expected_records,
            StreamName=expected_stream_name
        ).return_val(expected_put_records_batch_resp)

        # Test put_records()
        kw = KinesisWriter(aws_session=mock_session)
        resp = kw.put_records(
            stream_name=expected_stream_name,
            records=expected_records
        )
        assert resp == expected_resp
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_put_records_quit_on_err(self):
        # Set expected values
        expected_records = [{'Data': b'Some data', 'PartitionKey': 'partitionkey'}]
        expected_stream_name = 'kinesis-stream-name'
        expected_put_records_batch_resp = {
            'FailedRecordCount': 1,
            'Records': [{
                'ErrorMessage': 'Some error message'
            }]
        }

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'put_records',
            Records=expected_records,
            StreamName=expected_stream_name
        ).return_val(expected_put_records_batch_resp)

        # Test put_records() failure with quit_on_error=True
        kw = KinesisWriter(aws_session=mock_session)
        with pytest.raises(AssertionError) as e_info:
            kw.put_records(
                stream_name=expected_stream_name,
                records=expected_records,
                quit_on_error=True
            )
        assert str(e_info.value).startswith('1 record(s) failed to write')
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_put_records_no_quit_on_err(self):
        # Set expected values
        expected_records = [
            {'Data': b'Some data', 'PartitionKey': f'partitionkey{i}'}
            for i in range(MAX_PUT_RECORDS_BATCH * 3)
        ]
        expected_stream_name = 'kinesis-stream-name'
        fst_expected_err = Exception('first-test-exception')
        snd_expected_err = Exception('second-test-exception')
        thd_expected_err = Exception('third-test-exception')
        expected_exception = Exception([fst_expected_err, snd_expected_err, thd_expected_err])

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'put_records',
            Records=expected_records[:MAX_PUT_RECORDS_BATCH],
            StreamName=expected_stream_name
        ).raise_exception(fst_expected_err)
        mock_kinesis_client.on(
            'put_records',
            Records=expected_records[MAX_PUT_RECORDS_BATCH:-MAX_PUT_RECORDS_BATCH],
            StreamName=expected_stream_name
        ).raise_exception(snd_expected_err)
        mock_kinesis_client.on(
            'put_records',
            Records=expected_records[-MAX_PUT_RECORDS_BATCH:],
            StreamName=expected_stream_name
        ).raise_exception(thd_expected_err)

        # Test put_records() failure Exceptions quit_on_error=False
        kw = KinesisWriter(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            kw.put_records(
                stream_name=expected_stream_name,
                records=expected_records,
                quit_on_error=False
            )
        assert str(e_info.value) == str(expected_exception)
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_put_records_error(self):
        # Set expected values
        expected_records = [{'Data': b'Some data', 'PartitionKey': 'partitionkey'}]
        expected_stream_name = 'kinesis-stream-name'
        expected_err = Exception('test-exception')

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'put_records',
            Records=expected_records,
            StreamName=expected_stream_name
        ).raise_exception(expected_err)

        # Test put_records() Exception with quit_on_error=True
        kw = KinesisWriter(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            kw.put_records(
                stream_name=expected_stream_name,
                records=expected_records,
                quit_on_error=True
            )
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__put_records_batch(self):
        # Set expected values
        expected_batch = [{'a': 1}, {'b': 2}, {'c': 3}]
        expected_stream_name = 'kinesis-stream-name'
        expected_resp = {'Records': []}

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'put_records',
            Records=expected_batch,
            StreamName=expected_stream_name
        ).return_val(expected_resp)

        # Test __put_records_batch()
        kw = KinesisWriter(aws_session=mock_session)
        resp = kw._KinesisWriter__put_records_batch(
            batch=expected_batch,
            stream_name=expected_stream_name
        )
        assert resp == expected_resp
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__put_records_batch_throughput_exception(self):
        # Set expected values
        expected_batch = [{'a': 1}, {'b': 2}, {'c': 3}]
        expected_stream_name = 'kinesis-stream-name'
        expected_resp = {
            'FailedRecordCount': 1,
            'Records': [{'ErrorCode': PROVISIONED_THROUGHPUT_EXCEPTION}]
        }

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'put_records',
            Records=expected_batch,
            StreamName=expected_stream_name
        ).return_val(expected_resp)

        # Test __put_records_batch() ProvisionedThroughputExceededException
        kw = KinesisWriter(
            aws_session=mock_session,
            max_retry_count=1
        )
        with pytest.raises(AssertionError) as e_info:
            kw._KinesisWriter__put_records_batch(
                batch=expected_batch,
                stream_name=expected_stream_name
            )
        assert str(e_info.value) == 'Retries have been exhausted'
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__put_records_batch_failure(self):
        # Set expected values
        expected_batch = [{'a': 1}, {'b': 2}, {'c': 3}]
        expected_stream_name = 'kinesis-stream-name'
        expected_resp = {
            'FailedRecordCount': 1,
            'Records': [{'ErrorCode': 'SomeOtherFailure'}]
        }

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'put_records',
            Records=expected_batch,
            StreamName=expected_stream_name
        ).return_val(expected_resp)

        # Test __put_records_batch() failures
        kw = KinesisWriter(
            aws_session=mock_session,
            max_retry_count=1
        )
        resp = kw._KinesisWriter__put_records_batch(
            batch=expected_batch,
            stream_name=expected_stream_name
        )
        assert resp == expected_resp
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__put_records_batch_error(self):
        # Set expected values
        expected_batch = [{'a': 1}, {'b': 2}, {'c': 3}]
        expected_stream_name = 'kinesis-stream-name'
        expected_err = Exception('test-exception')

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'put_records',
            Records=expected_batch,
            StreamName=expected_stream_name
        ).raise_exception(expected_err)

        # Test __put_records_batch() Exception
        kw = KinesisWriter(
            aws_session=mock_session,
            max_retry_count=1
        )
        with pytest.raises(Exception) as e_info:
            kw._KinesisWriter__put_records_batch(
                batch=expected_batch,
                stream_name=expected_stream_name
            )
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__partition_put_records(self):
        mock_session, _ = self.get_mocks()
        kw = KinesisWriter(aws_session=mock_session)

        records_with_partition = kw._KinesisWriter__partition_put_records(
            records=[{'Data': b'Some data'}]
        )
        assert len(records_with_partition) == 1
        assert 'PartitionKey' in records_with_partition[0]
        assert isinstance(records_with_partition[0]['PartitionKey'], str)
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__batch_put_records(self):
        mock_session, _ = self.get_mocks()
        kw = KinesisWriter(aws_session=mock_session)

        batched = kw._KinesisWriter__batch_put_records(
            [{'a': 1}, {'b': 1}, {'c': 1}],
            batch_size=2
        )
        assert list(batched) == [[{'a': 1}, {'b': 1}], [{'c': 1}]]
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__validate_put_records(self):
        # Generate records to validate
        invalid_records = [
            {
                'Data': b'Some data',
                'PartitionKey': uuid.uuid4().hex,
                'ExplicitHashKey': uuid.uuid4().hex
            }
            for _ in range(MAX_INVALID_RECORDS + 2)
        ]
        valid_records = [
            {'Data': b'Some data'}
            for _ in range(5)
        ]
        records = valid_records + invalid_records

        # Set expected values
        expected_invalid_records = [
            {
                'record': r,
                'invalid_reason': 'Expected either record.PartitionKey or record.ExplicitHashKey, not both'
            }
            for r in invalid_records[:-2]
        ]
        
        # Test __validate_put_records()
        mock_session, _ = self.get_mocks()
        kw = KinesisWriter(aws_session=mock_session)
        actual_invalid_records = kw._KinesisWriter__validate_put_records(records)
        assert actual_invalid_records == expected_invalid_records
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__validate_put_record(self):
        mock_session, _ = self.get_mocks()
        kw = KinesisWriter(aws_session=mock_session)

        # Test has PartitionKey
        valid, invalid_reason = kw._KinesisWriter__validate_put_record({
            'Data': b'Some data',
            'PartitionKey': 'partitionkey'
        })
        assert valid
        assert not invalid_reason

        # Test has ExplicitHashKey
        valid, invalid_reason = kw._KinesisWriter__validate_put_record({
            'Data': b'Some data',
            'ExplicitHashKey': 'explicithashkey'
        })
        assert valid
        assert not invalid_reason

        # Test invalid record
        valid, invalid_reason = kw._KinesisWriter__validate_put_record({
            'Data': b'Some data',
            'PartitionKey': 'partitionkey',
            'ExplicitHashKey': 'explicithashkey'
        })
        assert not valid
        assert invalid_reason == 'Expected either record.PartitionKey or record.ExplicitHashKey, not both'

    @pytest.mark.unit
    def test__put_records_resp_has_throughput_exception(self):
        mock_session, _ = self.get_mocks()
        kw = KinesisWriter(aws_session=mock_session)
        resp = kw._KinesisWriter__put_records_resp_has_throughput_exception({
            'FailedRecordCount': 2,
            'Records': [{
                'ErrorCode': PROVISIONED_THROUGHPUT_EXCEPTION
            }, {
                'ErrorCode': 'SomeOtherError'
            }]
        })
        has_throughput_exception, has_other_exception = resp[0], resp[1]
        assert has_throughput_exception
        assert has_other_exception
        mock_session.assert_expectations()

    def get_mocks(self) -> Tuple[MockApi, MockApi]:
        mock_kinesis_client = MockApi()
        mock_session = MockApi()
        mock_session.on('client', 'kinesis').return_val(mock_kinesis_client)
        return mock_session, mock_kinesis_client
