import pytest
import pytz

from botocore.exceptions import ClientError
from datetime import date, datetime
from test_utils.mock import MockApi
from timestream.api import TimestreamLoader as ApiTimestreamLoader
from timestream.timestream_reader import DEFAULT_WRITE_CLIENT_CONFIG
from timestream.timestream_loader import TimestreamLoader
from timestream.dimensions import TimestreamDimensions
from timestream.load_result import TimestreamLoadResult
from timestream.record import RejectedTimestreamRecord, TimestreamRecord


class TestTimestreamLoader_UnitTests:
    @pytest.mark.unit
    def test_load(self):
        mock_session = self.get_mock_session()
        loader = TimestreamLoader("test.db", "test.table", mock_session)
        mock_session.assert_expectations()

        result = loader.load(data=[], capture_time=datetime(2021, 1, 1, 0, 0, 0, 0, pytz.UTC))
        assert type(result) == dict
        missing_keys = [
            k
            for k in ["error", "invalid_items", "accepted_records", "rejected_records"]
            if k not in result
        ]
        assert not missing_keys, f"Result {result} is missing keys {missing_keys}"
        assert type(result["invalid_items"]) == list
        assert len(result["invalid_items"]) == 0
        assert type(result["accepted_records"]) == list
        assert len(result["accepted_records"]) == 0
        assert type(result["rejected_records"]) == list
        assert len(result["rejected_records"]) == 0
        assert result["error"] == None
        assert loader.database_name == "test.db"
        assert loader.table_name == "test.table"

    @pytest.mark.unit
    def test_load_preview(self):
        mock_session = self.get_mock_session()
        loader = TimestreamLoader("test", "test", mock_session)
        mock_session.assert_expectations()

        result = loader.load_preview(
            data=[], capture_time=datetime(2021, 1, 1, 0, 0, 0, 0, pytz.UTC)
        )
        assert type(result) == dict
        missing_keys = [
            k
            for k in ["error", "invalid_items", "accepted_records", "rejected_records"]
            if k not in result
        ]
        assert not missing_keys, f"Result {result} is missing keys {missing_keys}"
        assert type(result["invalid_items"]) == list
        assert len(result["invalid_items"]) == 0
        assert type(result["accepted_records"]) == list
        assert len(result["accepted_records"]) == 0
        assert type(result["rejected_records"]) == list
        assert len(result["rejected_records"]) == 0
        assert result["error"] == None

    @pytest.mark.unit
    def test_private_load_no_data(self):
        mock_session = self.get_mock_session()
        loader = TimestreamLoader("test", "test", mock_session)
        mock_session.assert_expectations()

        result = loader._TimestreamLoader__load(
            data=[], capture_time=datetime(2021, 1, 1, 0, 0, 0, 0, pytz.UTC), preview=True
        )
        assert type(result) == TimestreamLoadResult
        assert len(result.invalid_items) == 0
        assert len(result.accepted_records) == 0
        assert len(result.rejected_records) == 0
        assert result.error == None

    @pytest.mark.unit
    def test_private_load_insert(self):
        # Test insert
        expected_source = "test soure"
        expected_reason = "test reason"
        expected_error = ClientError(
            {"RejectedRecords": [{"RecordIndex": 1, "Reason": expected_reason}]}, "test operation"
        )
        expected_records = [
            {"MeasureName": "valid", "MeasureValue": "1.0", "Dimensions": []},
            {"MeasureName": "invalid", "MeasureValue": "0.0", "Dimensions": []},
            {
                "MeasureName": "valid",
                "MeasureValue": "2.0",
                "Dimensions": [
                    {
                        "Name": "notshared",
                        "Value": "not_shared_dimension",
                        "DimensionValueType": "VARCHAR",
                    }
                ],
            },
        ]
        expected_common_dimensions = [
            {"Name": "shared", "Value": "shared_dimension", "DimensionValueType": "VARCHAR"},
            {"Name": "source", "Value": "test soure", "DimensionValueType": "VARCHAR"},
        ]
        expected_capture_time = "1609459200000"

        mock_client = MockApi()
        mock_client.on(
            "write_records",
            DatabaseName="test",
            TableName="test",
            CommonAttributes={
                "Dimensions": expected_common_dimensions,
                "MeasureValueType": "DOUBLE",
                "Time": expected_capture_time,
                "Version": 1,
            },
            Records=expected_records,
        ).raise_exception(expected_error)
        mock_session = self.get_mock_session(mock_client)

        loader = TimestreamLoader("test", "test", mock_session)
        result = loader._TimestreamLoader__load(
            data=[
                {"valid": 1.0, "invalid": 0.0, "shared": "shared_dimension"},
                {"valid": 2.0, "shared": "shared_dimension", "notshared": "not_shared_dimension"},
            ],
            capture_time=datetime(2021, 1, 1, 0, 0, 0, 0, pytz.UTC),
            source=expected_source,
            is_multi=False,
        )
        assert type(result) == TimestreamLoadResult
        assert len(result.accepted_records) == 2
        assert result.accepted_records[0].measure_name == "valid"
        assert len(result.rejected_records) == 1
        rejected_record = result.rejected_records[0]
        assert type(rejected_record) == RejectedTimestreamRecord
        assert str(rejected_record.error) == expected_reason
        assert rejected_record.record.measure_name == "invalid"
        mock_session.assert_expectations()
        mock_client.assert_expectations()

    @pytest.mark.unit
    def test_private_load_upsert(self):
        # Test upsert
        expected_capture_time = "1609459200000"
        expected_records = [
            {
                "MeasureName": "valid",
                "MeasureValue": "1.0",
                "Dimensions": [
                    {"Name": "test_dimension", "Value": "test", "DimensionValueType": "VARCHAR"}
                ],
                "Time": expected_capture_time,
            }
        ]

        mock_client = MockApi()
        mock_client.on(
            "write_records",
            DatabaseName="test",
            TableName="test",
            CommonAttributes={
                "Dimensions": [],
                "MeasureValueType": "DOUBLE",
                "Version": int(expected_capture_time),
            },
            Records=expected_records,
        ).return_val(None)
        mock_session = self.get_mock_session(mock_client)

        loader = TimestreamLoader(
            "test", "test", mock_session, current_time_ms=int(expected_capture_time)
        )
        result = loader._TimestreamLoader__load(
            data=[{"valid": 1.0, "time": int(expected_capture_time), "test_dimension": "test"}],
            capture_time=datetime(2021, 1, 1, 0, 0, 0, 0, pytz.UTC),
            upsert=True,
            is_multi=False,
        )
        assert type(result) == TimestreamLoadResult
        assert len(result.accepted_records) == 1
        assert result.accepted_records[0].measure_name == "valid"
        assert len(result.rejected_records) == 0
        mock_session.assert_expectations()
        mock_client.assert_expectations()

    @pytest.mark.unit
    def test_private_load_empty_dimension(self):
        # Test insert with empty dimension
        expected_capture_time = "1609459200000"
        expected_records = [
            {
                "MeasureName": "valid",
                "MeasureValue": "1.0",
                "Dimensions": [],
                "Time": expected_capture_time,
            }
        ]

        mock_client = MockApi()
        mock_client.on(
            "write_records",
            DatabaseName="test",
            TableName="test",
            CommonAttributes={
                "Dimensions": [],
                "MeasureValueType": "DOUBLE",
                "Version": int(expected_capture_time),
            },
            Records=expected_records,
        ).return_val(None)
        mock_session = self.get_mock_session(mock_client)

        loader = TimestreamLoader(
            "test", "test", mock_session, current_time_ms=int(expected_capture_time)
        )
        result = loader._TimestreamLoader__load(
            data=[{"valid": 1.0, "empty_dimension": "", "time": int(expected_capture_time)}],
            capture_time=datetime(2021, 1, 1, 0, 0, 0, 0, pytz.UTC),
            upsert=True,
            is_multi=False,
        )
        assert type(result) == TimestreamLoadResult
        assert len(result.accepted_records) == 1
        assert len(result.rejected_records) == 0
        mock_session.assert_expectations()
        mock_client.assert_expectations()

    @pytest.mark.unit
    def test_write_insert_records(self):
        # Set expectations
        expected_reason = "test reason"
        expected_error = ClientError(
            {"RejectedRecords": [{"RecordIndex": 1, "Reason": expected_reason}]}, "test operation"
        )
        expected_valid_record = TimestreamRecord("valid", 1.0)
        expected_invalid_record = TimestreamRecord("invalid", 0.0)
        expected_records = [
            {"MeasureName": "valid", "MeasureValue": "1.0", "Dimensions": []},
            {"MeasureName": "invalid", "MeasureValue": "0.0", "Dimensions": []},
        ]
        expected_capture_time = "1609459200000"
        expected_source = "test"

        # Set up mocks
        mock_client = MockApi()
        mock_client.on(
            "write_records",
            DatabaseName="test",
            TableName="test",
            CommonAttributes={
                "Dimensions": [],
                "MeasureValueType": "DOUBLE",
                "Time": expected_capture_time,
                "Version": 1,
            },
            Records=expected_records,
        ).raise_exception(expected_error)
        mock_session = self.get_mock_session(mock_client)

        # Test ClientError occurs during Timestream write
        loader = TimestreamLoader("test", "test", mock_session)
        accepted_records, rejected_records = loader._TimestreamLoader__write_insert_records(
            common_dimensions=TimestreamDimensions(),
            records=[expected_valid_record, expected_invalid_record],
            source=expected_source,
            capture_time=datetime(2021, 1, 1, 0, 0, 0, 0, tzinfo=pytz.UTC),
            is_multi=False,
        )
        assert len(accepted_records) == 1
        assert accepted_records[0] == expected_valid_record
        assert len(rejected_records) == 1

        rejected_record = rejected_records[0]
        assert type(rejected_record) == RejectedTimestreamRecord
        assert str(rejected_record.error) == expected_reason
        assert rejected_record.record == expected_invalid_record

        mock_session.assert_expectations()
        mock_client.assert_expectations()

    @pytest.mark.unit
    def test_write_upsert_records(self):
        # Set expectations
        expected_version = 1609459200000
        expected_reason = "test reason"
        expected_error = ClientError(
            {"RejectedRecords": [{"RecordIndex": 1, "Reason": expected_reason}]}, "test operation"
        )
        expected_valid_record = TimestreamRecord("valid", 1.0)
        expected_invalid_record = TimestreamRecord("invalid", 0.0)
        expected_records = [
            {"MeasureName": "valid", "MeasureValue": "1.0", "Dimensions": []},
            {"MeasureName": "invalid", "MeasureValue": "0.0", "Dimensions": []},
        ]
        expected_source = "test"

        # Set up mocks
        mock_client = MockApi()
        mock_client.on(
            "write_records",
            DatabaseName="test",
            TableName="test",
            CommonAttributes={
                "Dimensions": [],
                "MeasureValueType": "DOUBLE",
                "Version": expected_version,
            },
            Records=expected_records,
        ).raise_exception(expected_error)
        mock_session = self.get_mock_session(mock_client)

        # Test ClientError occurs during Timestream write
        loader = TimestreamLoader("test", "test", mock_session, current_time_ms=expected_version)
        accepted_records, rejected_records = loader._TimestreamLoader__write_upsert_records(
            common_dimensions=TimestreamDimensions(),
            records=[expected_valid_record, expected_invalid_record],
            source=expected_source,
            is_multi=False,
        )
        assert len(accepted_records) == 1
        assert accepted_records[0] == expected_valid_record
        assert len(rejected_records) == 1

        rejected_record = rejected_records[0]
        assert type(rejected_record) == RejectedTimestreamRecord
        assert str(rejected_record.error) == expected_reason
        assert rejected_record.record == expected_invalid_record

        mock_session.assert_expectations()
        mock_client.assert_expectations()

    @pytest.mark.unit
    def test_process_client_error(self):
        mock_session = self.get_mock_session()
        loader = TimestreamLoader("test", "test", mock_session)
        mock_session.assert_expectations()

        expected_reason = "test reason"
        expected_error = ClientError(
            {"RejectedRecords": [{"RecordIndex": 1, "Reason": expected_reason}]}, "test operation"
        )
        expected_valid_record = TimestreamRecord("valid", 1.0)
        expected_invalid_record = TimestreamRecord("invalid", 0.0)

        accepted_records, rejected_records = loader._TimestreamLoader__process_client_error(
            e=expected_error, records=[expected_valid_record, expected_invalid_record]
        )
        assert len(accepted_records) == 1
        assert accepted_records[0] == expected_valid_record
        assert len(rejected_records) == 1

        rejected_record = rejected_records[0]
        assert type(rejected_record) == RejectedTimestreamRecord
        assert str(rejected_record.error) == expected_reason
        assert rejected_record.record == expected_invalid_record

    @pytest.mark.unit
    def test_parse_data(self):
        mock_session = self.get_mock_session()
        loader = TimestreamLoader("test", "test", mock_session)

        # Test no data
        result = loader._TimestreamLoader__parse_data([], False, False, False)
        assert type(result) == TimestreamLoadResult
        assert str(result)
        assert len(result.invalid_items) == 0
        assert len(result.accepted_records) == 0
        assert len(result.rejected_records) == 0
        assert result.error == None

        # Test mix of valid and invalid items
        expected_valid_item = {"integer": 10, "dimension": "str"}
        expected_invalid_item = {"unparseable": datetime(2021, 1, 1, 0, 0, 0, 0, tzinfo=pytz.UTC)}
        expected_no_dimension_item = {"integer": 5}
        expected_error_msg = "Some items could not be parsed into Timestream records"

        result = loader._TimestreamLoader__parse_data(
            data=[
                expected_valid_item,
                expected_invalid_item,
                expected_no_dimension_item,
            ],
            upsert=False,
            ignore_parse_error=False,
            is_multi=False,
        )
        assert type(result) == TimestreamLoadResult
        assert len(result.invalid_items) == 2
        assert len(result.accepted_records) == 1
        assert len(result.rejected_records) == 0
        assert str(result.error) == expected_error_msg
        assert result.invalid_items[0]["item"] == expected_invalid_item
        assert result.invalid_items[1]["item"] == expected_no_dimension_item

        record = result.accepted_records[0]
        assert record and type(record) == TimestreamRecord
        assert record.measure_name == "integer"
        assert record.measure_value == 10.0
        assert record.mstime == None
        assert record.time == None

    @pytest.mark.unit
    def test_parse_item(self):
        mock_session = self.get_mock_session()
        loader = TimestreamLoader("test", "test", mock_session)
        mock_session.assert_expectations()

        # Test no item
        assert loader._TimestreamLoader__parse_item({}, False, False, False) == []

        # Test parse insert item
        expected_dimensions = TimestreamDimensions(
            {"string": "dimension", "remove this": "dimension"}
        )
        expected_dimensions.remove_dimension("remove this")
        records = loader._TimestreamLoader__parse_item(
            {"integer": 10, "float": 20.0, "list": [10, 20.0], "string": "dimension", "none": None},
            False,
            False,
            False,
        )
        assert len(records) == 3
        assert str(records[0])

        record = records[0]
        assert record and type(record) == TimestreamRecord
        assert record.measure_name == "integer"
        assert record.measure_value == 10.0
        assert record.mstime == None
        assert record.time == None
        assert record.dimensions == expected_dimensions
        assert record.dimensions.get_dimension("string") == "dimension"
        assert str(record.dimensions)

        record = records[1]
        assert record and type(record) == TimestreamRecord
        assert record.measure_name == "float"
        assert record.measure_value == 20.0
        assert record.mstime == None
        assert record.time == None
        assert record.dimensions == expected_dimensions

        record = records[2]
        assert record and type(record) == TimestreamRecord
        assert record.measure_name == "list"
        assert record.measure_value == 2.0
        assert record.mstime == None
        assert record.time == None
        assert record.dimensions == expected_dimensions

        # Test parse upsert item
        records = loader._TimestreamLoader__parse_item(
            {"Time": 1609459200000, "integer": 10, "string": "dimension"}, True, False, False
        )
        assert len(records) == 1

        record = records[0]
        assert record and type(record) == TimestreamRecord
        assert record.measure_name == "integer"
        assert record.measure_value == 10.0
        assert record.mstime == 1609459200000
        assert record.time == datetime.combine(date(2021, 1, 1), datetime.min.time()).replace(
            tzinfo=pytz.UTC
        )

        # Test parse error
        with pytest.raises(AssertionError) as e_info:
            loader._TimestreamLoader__parse_item(
                {"unparseable": {"key": "value"}, "string": "dimension"}, False, False, False
            )
        assert (
            str(e_info.value)
            == "Could not create a record or dimension from type \"<class 'dict'>\" at key unparseable"
        )

        # Test ignore parse error
        records = loader._TimestreamLoader__parse_item(
            {"unparseable": {"key": "value"}}, False, True, False
        )
        assert len(records) == 0

    @pytest.mark.unit
    def test_parse_item_empty_dimension(self):
        mock_session = self.get_mock_session()
        loader = TimestreamLoader("test", "test", mock_session)
        mock_session.assert_expectations()

        # Test parse insert item
        expected_dimensions = TimestreamDimensions({"string": "dimension"})
        records = loader._TimestreamLoader__parse_item(
            {
                "integer": 10,
                "float": 20.0,
                "list": [10, 20.0],
                "string": "dimension",
                "empty_dimension": "",
                "none": None,
            },
            False,
            False,
            False,
        )
        assert len(records) == 3
        assert str(records[0])

        record = records[0]
        assert record and type(record) == TimestreamRecord
        assert record.measure_name == "integer"
        assert record.measure_value == 10.0
        assert record.mstime == None
        assert record.time == None
        assert record.dimensions == expected_dimensions
        assert record.dimensions.get_dimension("string") == "dimension"
        assert str(record.dimensions)

        record = records[1]
        assert record and type(record) == TimestreamRecord
        assert record.measure_name == "float"
        assert record.measure_value == 20.0
        assert record.mstime == None
        assert record.time == None
        assert record.dimensions == expected_dimensions

        record = records[2]
        assert record and type(record) == TimestreamRecord
        assert record.measure_name == "list"
        assert record.measure_value == 2.0
        assert record.mstime == None
        assert record.time == None
        assert record.dimensions == expected_dimensions

    @pytest.mark.unit
    def test_does_database_exist(self):
        # Set expectations
        mock_response = {
            "Database": {
                "Arn": "arn:aws:timestream:us-east-1:000000000000:database/test",
                "DatabaseName": "test",
                "TableCount": 1,
                "KmsKeyId": "arn:aws:kms:us-east-1:000000000000:key/12345678-1234-1234-1234-1234567890ab",
                "CreationTime": datetime(2021, 7, 19, 16, 25, 26, 419000),
                "LastUpdatedTime": datetime(2021, 10, 8, 16, 25, 37, 893000),
            },
            "ResponseMetadata": {
                "RequestId": "1234567890abcdefghijklmnop",
                "HTTPStatusCode": 200,
                "HTTPHeaders": {
                    "x-amzn-requestid": "1234567890abcdefghijklmnop",
                    "content-type": "application/x-amz-json-1.0",
                    "content-length": "289",
                    "date": "Wed, 20 Oct 2021 16:54:48 GMT",
                },
                "RetryAttempts": 0,
            },
        }

        # Set up mocks
        mock_client = MockApi()
        mock_client.on(
            "describe_database",
            DatabaseName="test",
        ).return_val(mock_response)
        mock_session = self.get_mock_session(mock_client)

        # Test True when using TimestreamLoader's database attribute
        loader = TimestreamLoader("test", "test", mock_session)
        result = loader.does_database_exist()
        assert result
        mock_session.assert_expectations()
        mock_client.assert_expectations()

    @pytest.mark.unit
    def test_does_table_exist(self):
        # Set expectations
        mock_response = {
            "Table": {
                "Arn": "arn:aws:timestream:us-east-1:000000000000:database/test/table/test",
                "TableName": "test",
                "DatabaseName": "test",
                "TableStatus": "ACTIVE",
                "RetentionProperties": {
                    "MemoryStoreRetentionPeriodInHours": 720,
                    "MagneticStoreRetentionPeriodInDays": 1,
                },
                "CreationTime": datetime(2021, 8, 19, 18, 21, 1, 519000),
                "LastUpdatedTime": datetime(2021, 8, 31, 12, 36, 4, 503000),
            },
            "ResponseMetadata": {
                "RequestId": "SDKOJJS7STGCQXOFVLXPBYDRKI",
                "HTTPStatusCode": 200,
                "HTTPHeaders": {
                    "x-amzn-requestid": "SDKOJJS7STGCQXOFVLXPBYDRKI",
                    "content-type": "application/x-amz-json-1.0",
                    "content-length": "366",
                    "date": "Wed, 20 Oct 2021 17:41:09 GMT",
                },
                "RetryAttempts": 1,
            },
        }

        # Set up mocks
        mock_client = MockApi()
        mock_client.on("describe_table", DatabaseName="test", TableName="test").return_val(
            mock_response
        )
        mock_session = self.get_mock_session(mock_client)

        # Test True when using TimestreamLoader's database and table attribute
        loader = TimestreamLoader("test", "test", mock_session)
        result = loader.does_table_exist()
        assert result
        mock_session.assert_expectations()
        mock_client.assert_expectations()

    def get_mock_session(self, write_client=None) -> MockApi:
        mock_session = MockApi()
        mock_session.on(
            "client", "timestream-write", config=DEFAULT_WRITE_CLIENT_CONFIG
        ).return_val(write_client)
        mock_session.on("client", "timestream-query").return_val(MockApi())
        return mock_session
