import pytest
import pytz

from botocore.exceptions import ClientError
from datetime import datetime
from psi_kinesis.kinesis_reader import KinesisReader, PROVISIONED_THROUGHPUT_EXCEPTION
from psi_utils.unittest_mock import MockApi
from typing import Tuple


class TestKinesisReader_UnitTests:

    @pytest.mark.unit
    def test_assume_role(self):
        # Set expected values
        expected_role = 'RoleName'
        expected_account = '01234567890'
        expected_role_arn = 'arn:aws:iam::01234567890:role/RoleName'
        expected_role_session_name = 'RoleSession'
        expected_err = Exception('test-err')

        # Set up mocks
        mock_sts_client = MockApi()
        mock_session, _ = self.get_mocks()
        mock_session.on('client', 'sts').return_val(mock_sts_client)
        mock_sts_client.on(
            'assume_role',
            RoleArn=expected_role_arn,
            RoleSessionName=expected_role_session_name
        ).raise_exception(expected_err)

        # Test assume_role() Exception
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            kr.assume_role(
                expected_role,
                expected_account,
                role_session_name=expected_role_session_name
            )
        assert e_info.value == expected_err
        mock_sts_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_does_stream_exist(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_resp = {'StreamDescription': {'StreamName': expected_stream_name}}

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

        # Test does_stream_exist()
        kr = KinesisReader(aws_session=mock_session)
        assert kr.does_stream_exist(expected_stream_name)
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_does_stream_exist_resource_not_found(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_err = ClientError({
            'Error': {'Code': 'ResourceNotFoundException'}
        }, 'test')

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

        # Test does_stream_exist() ResourceNotFoundException
        kr = KinesisReader(aws_session=mock_session)
        assert not kr.does_stream_exist(expected_stream_name)
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_does_stream_exist_client_error(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_err = ClientError({
            'Error': {'Code': 'ClientError'}
        }, 'test')

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

        # Test does_stream_exist() ClientError
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(ClientError) as e_info:
            kr.does_stream_exist(expected_stream_name)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_does_stream_exist_error(self):
        # Set expected values
        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(
            'describe_stream',
            StreamName=expected_stream_name
        ).raise_exception(expected_err)

        # Test does_stream_exist() Exception
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            kr.does_stream_exist(expected_stream_name)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_list_shards(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_shard_filter = {
            'Type': 'FROM_TIMESTAMP',
            'Timestamp': datetime(2022, 4, 1, tzinfo=pytz.UTC)
        }
        expected_shards = [{
            'ShardId': 'shard-id-000001'
        }]

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'list_shards',
            StreamName=expected_stream_name,
            ShardFilter=expected_shard_filter
        ).return_val({'Shards': expected_shards})

        # Test list_shards()
        kr = KinesisReader(aws_session=mock_session)
        shards = kr.list_shards(
            expected_stream_name,
            shard_filter=expected_shard_filter
        )
        assert shards == expected_shards
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_list_shards_error(self):
        # Set expected values
        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(
            'list_shards',
            StreamName=expected_stream_name
        ).raise_exception(expected_err)

        # Test list_shards() Exception
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            kr.list_shards(expected_stream_name)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_shard_id = 'shard-id-000001'
        expected_shard_iterator = 'shard-iterator-id'
        expected_records = [{'a': 1}, {'b': 2}, {'c': 3}]

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_shard_iterator',
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType='TRIM_HORIZON'
        ).return_val({'ShardIterator': expected_shard_iterator})
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=expected_shard_iterator,
            Limit=10000
        ).return_val({
            'Records': expected_records,
            'MillisBehindLatest': 0
        })

        # Test get_records() with a sequence number but no shard iterator type
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(AssertionError) as e_info:
            kr.get_records(
                expected_stream_name,
                shard_id=expected_shard_id,
                starting_sequence_number='1234567890123456789'
            )
        assert str(e_info.value) == 'Cannot specify a starting sequence number or timestamp if shard iteratory type is one of [TRIM_HORIZON,LATEST]'

        # Test get_records() with a timestamp but no shard iterator type
        kr = KinesisReader(aws_session=mock_session)
        with pytest.raises(AssertionError) as e_info:
            kr.get_records(
                expected_stream_name,
                shard_id=expected_shard_id,
                timestamp=datetime(2022, 4, 1, tzinfo=pytz.UTC)
            )
        assert str(e_info.value) == 'Cannot specify a starting sequence number or timestamp if shard iteratory type is one of [TRIM_HORIZON,LATEST]'

        # Test get_records()
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records(
            expected_stream_name,
            shard_id=expected_shard_id
        )
        assert records == {expected_shard_id: expected_records}
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records_at_sequence_number(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_shard_id = 'shard-id-000001'
        expected_shard_iterator = 'shard-iterator-id'
        expected_shard_iterator_type = 'AT_SEQUENCE_NUMBER'
        expected_sequence_number = '1234567890123456789'
        expected_records = [{'a': 1}, {'b': 2}, {'c': 3}]

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_shard_iterator',
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            StartingSequenceNumber=expected_sequence_number
        ).return_val({'ShardIterator': expected_shard_iterator})
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=expected_shard_iterator,
            Limit=10000
        ).return_val({
            'Records': expected_records,
            'MillisBehindLatest': 0
        })

        # Test get_records() AT_SEQUENCE_NUMBER
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records(
            expected_stream_name,
            shard_id=expected_shard_id,
            shard_iterator_type=expected_shard_iterator_type,
            starting_sequence_number=expected_sequence_number
        )
        assert records == {expected_shard_id: expected_records}
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records_at_timestamp(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_shard_id = 'shard-id-000001'
        expected_shard_iterator = 'shard-iterator-id'
        expected_shard_iterator_type = 'AT_TIMESTAMP'
        expected_timestamp = datetime(2022, 4, 1, tzinfo=pytz.UTC)
        expected_records = [{'a': 1}, {'b': 2}, {'c': 3}]

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_shard_iterator',
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            Timestamp=expected_timestamp
        ).return_val({'ShardIterator': expected_shard_iterator})
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=expected_shard_iterator,
            Limit=10000
        ).return_val({
            'Records': expected_records,
            'MillisBehindLatest': 0
        })

        # Test get_records() AT_TIMESTAMP
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records(
            expected_stream_name,
            shard_id=expected_shard_id,
            shard_iterator_type=expected_shard_iterator_type,
            timestamp=expected_timestamp
        )
        assert records == {expected_shard_id: expected_records}
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records_for_entire_stream(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_shard_iterator = 'shard-iterator-id'
        expected_records = [{'a': 1}, {'b': 2}, {'c': 3}]
        expected_shard_id = 'shard-id-000001'

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'list_shards',
            StreamName=expected_stream_name
        ).return_val({'Shards': [{'ShardId': expected_shard_id}]})
        mock_kinesis_client.on(
            'get_shard_iterator',
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType='TRIM_HORIZON'
        ).return_val({'ShardIterator': expected_shard_iterator})
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=expected_shard_iterator,
            Limit=10000
        ).return_val({
            'Records': expected_records,
            'MillisBehindLatest': 0
        })

        # Test get_records() for an entire Kinesis stream
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records(expected_stream_name)
        assert records == {expected_shard_id: expected_records}
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_get_records_with_shard_iterator(self):
        # Set expected values
        fst_shard_iterator = 'first-shard-iterator-id'
        snd_shard_iterator = 'second-shard-iterator-id'
        fst_records = [{'a': 1}, {'b': 2}]
        snd_records = [{'c': 3}, {'d': 4}]
        fst_resp = {
            'Records': fst_records,
            'NextShardIterator': snd_shard_iterator,
            'MillisBehindLatest': 123
        }
        snd_resp = {
            'Records': snd_records,
            'NextShardIterator': 'third-shard-iterator-id',
            'MillisBehindLatest': 0
        }
        expected_records = fst_records + snd_records

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=fst_shard_iterator,
            Limit=10000
        ).return_val(fst_resp)
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=snd_shard_iterator,
            Limit=10000
        ).return_val(snd_resp)

        # Test get_records_with_shard_iterator()
        kr = KinesisReader(aws_session=mock_session)
        records = kr.get_records_with_shard_iterator(fst_shard_iterator)
        assert records == expected_records
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_shard_iterator(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_shard_id = 'shard-id-000001'
        expected_shard_iterator_type = 'AT_SEQUENCE_NUMBER'
        expected_sequence_number = '1234567890123456789'
        expected_shard_iterator = 'shard-iterator-id'
        expected_resp = {
            'ShardIterator': expected_shard_iterator
        }

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_shard_iterator',
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            StartingSequenceNumber=expected_sequence_number
        ).return_val(expected_resp)

        # Test __get_shard_iterator()
        kr = KinesisReader(aws_session=mock_session)
        shard_iterator = kr._KinesisReader__get_shard_iterator(
            expected_stream_name,
            expected_shard_id,
            expected_shard_iterator_type,
            starting_sequence_number=expected_sequence_number
        )
        assert shard_iterator == expected_shard_iterator
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_shard_iterator_throughput_exception(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_shard_id = 'shard-id-000001'
        expected_shard_iterator_type = 'AT_TIMESTAMP'
        expected_timestamp = datetime(2022, 4, 1, tzinfo=pytz.UTC)
        expected_err = ClientError({
            'Error': {'Code': PROVISIONED_THROUGHPUT_EXCEPTION}
        }, 'test')

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_shard_iterator',
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            Timestamp=expected_timestamp
        ).raise_exception(expected_err)

        # Test __get_shard_iterator() ProvisionedThroughputExceededException
        kr = KinesisReader(
            aws_session=mock_session,
            max_retry_count=1
        )

        with pytest.raises(AssertionError) as e_info:
            kr._KinesisReader__get_shard_iterator(
                expected_stream_name,
                expected_shard_id,
                expected_shard_iterator_type,
                timestamp=expected_timestamp
            )
        assert str(e_info.value) == 'Retries have been exhausted'
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_shard_iterator_client_error(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_shard_id = 'shard-id-000001'
        expected_shard_iterator_type = 'AT_TIMESTAMP'
        expected_timestamp = datetime(2022, 4, 1, tzinfo=pytz.UTC)
        expected_err = ClientError({
            'Error': {'Code': 'ClientError'}
        }, 'test')

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_shard_iterator',
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            Timestamp=expected_timestamp
        ).raise_exception(expected_err)

        # Test __get_shard_iterator() ClientError
        kr = KinesisReader(
            aws_session=mock_session,
            max_retry_count=1
        )

        with pytest.raises(ClientError) as e_info:
            kr._KinesisReader__get_shard_iterator(
                expected_stream_name,
                expected_shard_id,
                expected_shard_iterator_type,
                timestamp=expected_timestamp
            )
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_shard_iterator_error(self):
        # Set expected values
        expected_stream_name = 'kinesis-stream-name'
        expected_shard_id = 'shard-id-000001'
        expected_shard_iterator_type = 'AT_TIMESTAMP'
        expected_timestamp = datetime(2022, 4, 1, tzinfo=pytz.UTC)
        expected_err = Exception('test-exception')

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_shard_iterator',
            StreamName=expected_stream_name,
            ShardId=expected_shard_id,
            ShardIteratorType=expected_shard_iterator_type,
            Timestamp=expected_timestamp
        ).raise_exception(expected_err)

        # Test __get_shard_iterator() Exception
        kr = KinesisReader(
            aws_session=mock_session,
            max_retry_count=1
        )

        with pytest.raises(Exception) as e_info:
            kr._KinesisReader__get_shard_iterator(
                expected_stream_name,
                expected_shard_id,
                expected_shard_iterator_type,
                timestamp=expected_timestamp
            )
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_records_with_shard_iterator(self):
        # Set expected values
        expected_shard_iterator = 'shard-iterator-id'
        expected_resp = {
            'Records': [],
            'NextShardIterator': None,
            'MillisBehindLatest': 123
        }

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=expected_shard_iterator,
            Limit=10000
        ).return_val(expected_resp)

        # Test __get_records_with_shard_iterator()
        kr = KinesisReader(aws_session=mock_session)
        resp = kr._KinesisReader__get_records_with_shard_iterator(expected_shard_iterator)
        assert resp == expected_resp
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_records_with_shard_iterator_throughput_exception(self):
        # Set expected values
        expected_shard_iterator = 'shard-iterator-id'
        expected_err = ClientError({
            'Error': {'Code': PROVISIONED_THROUGHPUT_EXCEPTION}
        }, 'test')

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=expected_shard_iterator,
            Limit=10000
        ).raise_exception(expected_err)

         # Test __get_records_with_shard_iterator() ProvisionedThroughputExceededException
        kr = KinesisReader(
            aws_session=mock_session,
            max_retry_count=1
        )

        with pytest.raises(AssertionError) as e_info:
            kr._KinesisReader__get_records_with_shard_iterator(expected_shard_iterator)
        assert str(e_info.value) == 'Retries have been exhausted'
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_records_with_shard_iterator_client_error(self):
        # Set expected values
        expected_shard_iterator = 'shard-iterator-id'
        expected_err = ClientError({
            'Error': {'Code': 'ClientError'}
        }, 'test')

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=expected_shard_iterator,
            Limit=10000
        ).raise_exception(expected_err)

         # Test __get_records_with_shard_iterator() ClientError
        kr = KinesisReader(
            aws_session=mock_session,
            max_retry_count=1
        )

        with pytest.raises(ClientError) as e_info:
            kr._KinesisReader__get_records_with_shard_iterator(expected_shard_iterator)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test__get_records_with_shard_iterator_error(self):
        # Set expected values
        expected_shard_iterator = 'shard-iterator-id'
        expected_err = Exception('test exception')

        # Set up mocks
        mock_session, mock_kinesis_client = self.get_mocks()
        mock_kinesis_client.on(
            'get_records',
            ShardIterator=expected_shard_iterator,
            Limit=10000
        ).raise_exception(expected_err)

         # Test __get_records_with_shard_iterator() Exception
        kr = KinesisReader(
            aws_session=mock_session,
            max_retry_count=1
        )

        with pytest.raises(Exception) as e_info:
            kr._KinesisReader__get_records_with_shard_iterator(expected_shard_iterator)
        assert e_info.value == expected_err
        mock_kinesis_client.assert_expectations()
        mock_session.assert_expectations()
    
    @pytest.mark.unit
    def test_next_wait_seconds(self):
        retry_backoff_rate = 2.5
        mock_session, _ = self.get_mocks()
        kr = KinesisReader(
            aws_session=mock_session,
            retry_backoff_rate=retry_backoff_rate
        )
        assert kr._next_wait_seconds(None) == 1.0
        assert kr._next_wait_seconds(1.0) == retry_backoff_rate
        assert kr._next_wait_seconds(retry_backoff_rate) == 6.25
        assert kr._next_wait_seconds(10) == 25

    @pytest.mark.unit
    def test_next_retry_count(self):
        mock_session, _ = self.get_mocks()
        kr = KinesisReader(aws_session=mock_session)
        assert kr._next_retry_count(None) == 1
        assert kr._next_retry_count(10) == 11

    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
