import pytest
import pytz

from botocore.exceptions import ClientError
from datetime import date, datetime
from timestream.timestream_reader import DEFAULT_WRITE_CLIENT_CONFIG, TimestreamReader
from test_utils.mock import MockApi


class TestTimestreamLoader_UnitTests:

    @pytest.mark.unit
    def test_does_database_exist(self):
        # Set expectations
        expected_error = ClientError(
            {
                "response": {
                    "Error": {
                        "Message": "The database default does not exist.",
                        "Code": "ResourceNotFoundException",
                    },
                    "ResponseMetadata": {
                        "RequestId": "V6E3FBONKV7C6YK3WLSVISIJGI",
                        "HTTPStatusCode": 404,
                        "HTTPHeaders": {
                            "x-amzn-requestid": "V6E3FBONKV7C6YK3WLSVISIJGI",
                            "content-type": "application/x-amz-json-1.0",
                            "content-length": "129",
                            "date": "Wed, 20 Oct 2021 17:04:42 GMT",
                        },
                        "RetryAttempts": 0,
                    },
                    "Message": "The database default does not exist.",
                },
                "operation_name": "DescribeDatabase",
            },
            "DescribeDatabase",
        )
        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_write_client = MockApi()
        mock_write_client.on(
            "describe_database",
            DatabaseName="default",
        ).raise_exception(expected_error)
        mock_write_client.on(
            "describe_database",
            DatabaseName="test",
        ).return_val(mock_response)
        mock_session = self.get_mock_session(write_client=mock_write_client)

        # Test False when using TimestreamReader's database attribute
        reader = TimestreamReader(mock_session)
        result = reader.does_database_exist("default")
        assert not result

        # Test exception raised if not found
        with pytest.raises(ClientError):
            reader.does_database_exist("default", raise_err=True)

        # Test True using test name
        result = reader.does_database_exist("test")
        assert result

        mock_session.assert_expectations()
        mock_write_client.assert_expectations()

    @pytest.mark.unit
    def test_does_table_exist(self):
        # Set expectations
        expected_error = ClientError(
            {
                "response": {
                    "Error": {
                        "Message": "The table default does not exist.",
                        "Code": "ResourceNotFoundException",
                    },
                    "ResponseMetadata": {
                        "RequestId": "V6E3FBONKV7C6YK3WLSVISIJGI",
                        "HTTPStatusCode": 404,
                        "HTTPHeaders": {
                            "x-amzn-requestid": "V6E3FBONKV7C6YK3WLSVISIJGI",
                            "content-type": "application/x-amz-json-1.0",
                            "content-length": "129",
                            "date": "Wed, 20 Oct 2021 17:04:42 GMT",
                        },
                        "RetryAttempts": 0,
                    },
                    "Message": "The table default does not exist.",
                },
                "operation_name": "DescribeTable",
            },
            "DescribeTable",
        )
        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_write_client = MockApi()
        mock_write_client.on(
            "describe_table", DatabaseName="default", TableName="default"
        ).raise_exception(expected_error)
        mock_write_client.on("describe_table", DatabaseName="test", TableName="test").return_val(
            mock_response
        )
        mock_session = self.get_mock_session(write_client=mock_write_client)

        # Test False when using TimestreamReader's database and table attribute
        reader = TimestreamReader(mock_session)
        result = reader.does_table_exist("default", "default")
        assert not result

        # Test exception raised if not found
        with pytest.raises(ClientError) as e_info:
            reader.does_table_exist("default", "default", raise_err=True)

        # Test True using test name
        result = reader.does_table_exist("test", "test")
        assert result

        mock_session.assert_expectations()
        mock_write_client.assert_expectations()

    @pytest.mark.unit
    def test_custom_read(self):
        # Set expectations
        expected_time = datetime(2022, 1, 1, 11, 12, 13, 14, pytz.utc)
        expected_time_str = expected_time.strftime("%Y-%m-%d %H:%M:%S.%f000")
        expected_inline_sql = 'SELECT * FROM "db"."table" LIMIT 5'
        expected_file_sql = 'SELECT *\nFROM "db"."table"\nLIMIT 5'
        expected_rows = [
            {
                "Data": [
                    {"ScalarValue": 1.0},
                    {"NullValue": True},
                    {"ScalarValue": expected_time_str},
                ]
            }
        ]
        expected_columns = [
            {"Name": "measure_value::double"},
            {"Name": "dimension1"},
            {"Name": "time"},
        ]
        expected_next_token = "token"
        expected_resp = {
            "Rows": expected_rows,
            "ColumnInfo": expected_columns,
            "NextToken": expected_next_token,
        }
        expected_records = [
            {"measure_value::double": 1.0, "dimension1": None, "time": expected_time}
        ]

        # Set up mocks
        mock_custom_read_client = MockApi()
        mock_custom_read_client.on(
            "query", QueryString=expected_inline_sql, MaxRows=1000
        ).return_val(expected_resp)
        mock_custom_read_client.on(
            "query", QueryString=expected_inline_sql, NextToken=expected_next_token, MaxRows=1000
        ).return_val({})
        mock_custom_read_client.on("query", QueryString=expected_file_sql, MaxRows=1000).return_val(
            expected_resp
        )
        mock_custom_read_client.on(
            "query", QueryString=expected_file_sql, NextToken=expected_next_token, MaxRows=1000
        ).return_val({})

        # Test custom_read with inline SQL
        reader = TimestreamReader(self.get_mock_session(read_client=mock_custom_read_client))
        records = reader.custom_read(sql='SELECT * FROM "db"."table" LIMIT 5')
        assert records == expected_records

        # Test custom_read with SQL from file
        reader = TimestreamReader(self.get_mock_session(read_client=mock_custom_read_client))
        records = reader.custom_read(file_path="test/queries/test-custom-read.sql")
        assert records == expected_records

    @pytest.mark.unit
    def test_read(self):
        # Set expectations
        expected_time = datetime(2022, 1, 1, 11, 12, 13, 14, pytz.utc)
        expected_time_str = expected_time.strftime("%Y-%m-%d %H:%M:%S.%f000")
        expected_sql = 'SELECT * FROM "db"."table" LIMIT 5'
        expected_rows = [
            {
                "Data": [
                    {"ScalarValue": 1.0},
                    {"NullValue": True},
                    {"ScalarValue": expected_time_str},
                ]
            }
        ]
        expected_columns = [
            {"Name": "measure_value::double"},
            {"Name": "dimension1"},
            {"Name": "time"},
        ]
        expected_next_token = "token"
        expected_resp = {
            "Rows": expected_rows,
            "ColumnInfo": expected_columns,
            "NextToken": expected_next_token,
        }
        expected_records = [
            {"measure_value::double": 1.0, "dimension1": None, "time": expected_time}
        ]

        # Set up mocks
        mock_read_client = MockApi()
        mock_read_client.on("query", QueryString=expected_sql, MaxRows=1000).return_val(
            expected_resp
        )
        mock_read_client.on(
            "query", QueryString=expected_sql, NextToken=expected_next_token, MaxRows=1000
        ).return_val({})

        # Test read
        reader = TimestreamReader(self.get_mock_session(read_client=mock_read_client))
        records = reader.read("db", "table", limit=5)
        assert records == expected_records

    @pytest.mark.unit
    def test_generate_sql(self):
        reader = TimestreamReader(self.get_mock_session())
        time = datetime(2022, 1, 1, 11, 12, 13, 14, pytz.utc)
        min_time = datetime.combine(time.date(), datetime.min.time()).replace(tzinfo=pytz.utc)
        max_time = datetime.combine(time.date(), datetime.max.time()).replace(tzinfo=pytz.utc)

        # Test projection expression
        sql = reader._generate_sql("db", "table", projection_expression="COUNT(*)")
        assert sql == 'SELECT COUNT(*) FROM "db"."table"'

        # Test measure value
        sql = reader._generate_sql("db", "table", measure_value=1.0, measure_value_op="<>")
        assert sql == 'SELECT * FROM "db"."table" WHERE measure_value::double <> 1.0'

        # Test time
        sql = reader._generate_sql("db", "table", time=time)
        assert (
            sql
            == 'SELECT * FROM "db"."table" WHERE time = FROM_ISO8601_TIMESTAMP(\'2022-01-01T11:12:13.000014+00:00\')'
        )

        # Test min time
        sql = reader._generate_sql("db", "table", min_time=min_time)
        assert (
            sql
            == 'SELECT * FROM "db"."table" WHERE time >= FROM_ISO8601_TIMESTAMP(\'2022-01-01T00:00:00+00:00\')'
        )

        # Test max time
        sql = reader._generate_sql("db", "table", max_time=max_time)
        assert (
            sql
            == 'SELECT * FROM "db"."table" WHERE time <= FROM_ISO8601_TIMESTAMP(\'2022-01-01T23:59:59.999999+00:00\')'
        )

        # Test min + max time
        sql = reader._generate_sql("db", "table", min_time=min_time, max_time=max_time)
        assert (
            sql
            == "SELECT * FROM \"db\".\"table\" WHERE time BETWEEN FROM_ISO8601_TIMESTAMP('2022-01-01T00:00:00+00:00') AND FROM_ISO8601_TIMESTAMP('2022-01-01T23:59:59.999999+00:00')"
        )

        # Test kwargs + limit
        sql = reader._generate_sql(
            "db", "table", limit=5, dim1="val1", t_AppID=["SVC98765", "SVC01234"]
        )
        assert (
            sql
            == "SELECT * FROM \"db\".\"table\" WHERE dim1 = 'val1' AND t_AppID IN ('SVC98765','SVC01234') LIMIT 5"
        )

    @pytest.mark.unit
    def test_generate_sql_assertions(self):
        reader = TimestreamReader(self.get_mock_session())
        time = datetime(2022, 1, 1, 11, 12, 13, 14, pytz.utc)
        min_time = datetime.combine(time.date(), datetime.min.time()).replace(tzinfo=pytz.utc)
        max_time = datetime.combine(time.date(), datetime.max.time()).replace(tzinfo=pytz.utc)

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql(None, "table")
        assert str(e_info.value).startswith("Cannot generate SQL without a database")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", None)
        assert str(e_info.value).startswith("Cannot generate SQL without a table")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", projection_expression=1)
        assert str(e_info.value).startswith("Invalid projection expression")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", measure_value="a")
        assert str(e_info.value).startswith("Invalid measure value")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", measure_value_op=10)
        assert str(e_info.value).startswith("Invalid measure value op")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", time=10)
        assert str(e_info.value).startswith("Invalid time")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", min_time=10)
        assert str(e_info.value).startswith("Invalid min_time")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", max_time=10)
        assert str(e_info.value).startswith("Invalid max_time")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", time=time, min_time=min_time)
        assert str(e_info.value).startswith("Cannot specify min_time when time is specified")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", time=time, max_time=max_time)
        assert str(e_info.value).startswith("Cannot specify max_time when time is specified")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", limit="10")
        assert str(e_info.value).startswith("Invalid limit")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", limit=-1)
        assert str(e_info.value).startswith("Invalid limit")

        with pytest.raises(AssertionError) as e_info:
            reader._generate_sql("db", "table", dim1=10)
        assert str(e_info.value).startswith("Invalid kwarg value")

    @pytest.mark.unit
    def test_parse_rows(self):
        # Set expectations
        expected_time = datetime(2022, 1, 1, 11, 12, 13, 14, pytz.utc)
        expected_time_str = expected_time.strftime("%Y-%m-%d %H:%M:%S.%f000")
        expected_rows = [
            {
                "Data": [
                    {"ScalarValue": 1.0},
                    {"NullValue": True},
                    {"ScalarValue": expected_time_str},
                ]
            },
            {
                "Data": [
                    {"ScalarValue": 2.0},
                    {"Unknown": "bad"},
                    {"ScalarValue": expected_time_str},
                ]
            },
        ]
        expected_columns = [
            {"Name": "measure_value::double"},
            {"Name": "dimension1"},
            {"Name": "time"},
        ]
        expected_records = [
            {"measure_value::double": 1.0, "dimension1": None, "time": expected_time}
        ]
        reader = TimestreamReader(self.get_mock_session())

        # Test with bad row
        with pytest.raises(Exception) as e_info:
            reader._parse_rows(expected_rows, expected_columns)
        assert str(e_info.value).startswith("Unrecognized type for data")

        # Test without bad row
        expected_rows = expected_rows[:1]
        records = reader._parse_rows(expected_rows, expected_columns)
        assert records == expected_records

    @pytest.mark.unit
    def test_parse_rows_assertions(self):
        reader = TimestreamReader(self.get_mock_session())

        with pytest.raises(AssertionError) as e_info:
            reader._parse_rows("rows", [])
        assert str(e_info.value).startswith("Invalid rows to parse")

        with pytest.raises(AssertionError) as e_info:
            reader._parse_rows([], "columns")
        assert str(e_info.value).startswith("Invalid columns to parse")

        with pytest.raises(AssertionError) as e_info:
            reader._parse_rows(["row"], [])
        assert str(e_info.value).startswith("Invalid row")

        with pytest.raises(AssertionError) as e_info:
            reader._parse_rows([{"Data": []}], [{"Name": "col1"}])
        assert str(e_info.value).startswith("Invalid row data")

        with pytest.raises(AssertionError) as e_info:
            reader._parse_rows([{"Data": ["val1"]}], ["column"])
        assert str(e_info.value).startswith("Invalid column")

        with pytest.raises(AssertionError) as e_info:
            reader._parse_rows([{"Data": ["val1"]}], [{"bad": "column"}])
        assert str(e_info.value).startswith("Invalid column")

    @pytest.mark.unit
    def test_dt_to_ms(self):
        mock_session = self.get_mock_session()
        reader = TimestreamReader(mock_session)
        mock_session.assert_expectations()

        dt = datetime.combine(date(2021, 1, 1), datetime.min.time()).replace(tzinfo=pytz.UTC)
        assert reader._dt_to_ms(dt) == 1609459200000
        assert reader._dt_to_ms(None) == -1

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