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
