import pytest

from io import BytesIO
from psi_athena.athena import Athena
from test_utils.mock import MockApi
from typing import Tuple


class TestAthena_UnitTests:

    @pytest.mark.unit
    def test_assume_role(self):
        # Set expected values
        expected_role = "aws:role:arn"
        expected_region = "region-code"
        expected_session_name = "test session name"
        expected_access_key = "accesskey"
        expected_secret_access_key = "secret"
        expected_session_token = "token"

        # Set up mocks
        mock_session, _ = self.get_mocks()
        mock_stsclient = MockApi()
        mock_stsclient.on(
            "assume_role", RoleArn=expected_role, RoleSessionName=expected_session_name
        ).return_val(
            {
                "Credentials": {
                    "AccessKeyId": expected_access_key,
                    "SecretAccessKey": expected_secret_access_key,
                    "SessionToken": expected_session_token,
                }
            }
        )
        mock_session.on(
            "client",
            "athena",
            region_name=expected_region,
            aws_access_key_id=expected_access_key,
            aws_secret_access_key=expected_secret_access_key,
            aws_session_token=expected_session_token,
        ).return_val(None)
        mock_session.on("client", "sts").return_val(mock_stsclient)

        # Test assume_role()
        a = Athena(aws_session=mock_session)
        a.assume_role(expected_role, expected_region, expected_session_name)
        mock_stsclient.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_run_athena_query_exception(self):
        # Set expected values
        expected_query = "sql"
        expected_database = "test_db"
        expected_table = "test_table"
        expected_catalog = "TestCatalog"
        expected_workgroup = "test-wg"
        expected_query_execution_id = "expected-id"
        expected_bucket = "bucket"
        expected_key = "key"
        expected_exception = Exception("test exception")

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            "start_query_execution",
            QueryString=expected_query,
            QueryExecutionContext={"Database": expected_database, "Catalog": expected_catalog},
            WorkGroup=expected_workgroup,
        ).return_val({"QueryExecutionId": expected_query_execution_id})
        mock_athena_client.on(
            "get_query_execution", QueryExecutionId=expected_query_execution_id
        ).return_val(
            {
                "QueryExecution": {
                    "Status": {"State": "SUCCEEDED"},
                    "ResultConfiguration": {
                        "OutputLocation": f"s3://{expected_bucket}/{expected_key}"
                    },
                }
            }
        )
        mock_s3_client = MockApi()
        mock_s3_client.on("get_object", Bucket=expected_bucket, Key=expected_key).raise_exception(
            expected_exception
        )
        mock_session.on("client", "s3").return_val(mock_s3_client)

        # Test run_athena_query() exception
        a = Athena(
            data_catalog=expected_catalog,
            workgroup=expected_workgroup,
            database=expected_database,
            aws_session=mock_session,
        )
        with pytest.raises(Exception) as e_info:
            a.run_athena_query(expected_query, table=expected_table, download_results=True)
        assert e_info.value == expected_exception
        mock_athena_client.assert_expectations()
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_start_query(self):
        # Set expected values
        expected_query = "sql"
        expected_database = "test_db"
        expected_catalog = "TestCatalog"
        expected_workgroup = "test-wg"
        expected_query_execution_id = "expected-id"

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            "start_query_execution",
            QueryString=expected_query,
            QueryExecutionContext={"Database": expected_database, "Catalog": expected_catalog},
            WorkGroup=expected_workgroup,
        ).return_val({"QueryExecutionId": expected_query_execution_id})

        # Test start_query()
        a = Athena(
            data_catalog=expected_catalog,
            workgroup=expected_workgroup,
            database=expected_database,
            aws_session=mock_session,
        )
        query_execution_id = a.start_query(expected_query)
        assert query_execution_id == expected_query_execution_id
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_start_query_exception(self):
        # Set expected values
        expected_query = "sql"
        expected_database = "test_db"
        expected_catalog = "TestCatalog"
        expected_workgroup = "test-wg"
        expected_exception = Exception("test")

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            "start_query_execution",
            QueryString=expected_query,
            QueryExecutionContext={"Database": expected_database, "Catalog": expected_catalog},
            WorkGroup=expected_workgroup,
        ).raise_exception(expected_exception)

        # Test start_query() exception
        a = Athena(
            data_catalog=expected_catalog,
            workgroup=expected_workgroup,
            database=expected_database,
            aws_session=mock_session,
        )
        with pytest.raises(Exception) as e_info:
            a.start_query(expected_query)
        assert e_info.value == expected_exception
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_poll_for_query_results(self):
        # Set expected values
        expected_query_execution_id = "expected-id"
        expected_results = {"QueryExecution": {"Status": {"State": "SUCCEEDED"}}}

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            "get_query_execution", QueryExecutionId=expected_query_execution_id
        ).return_val(expected_results)

        # Test poll_for_query_results()
        a = Athena(aws_session=mock_session)
        results = a.poll_for_query_results(expected_query_execution_id)
        assert results == expected_results
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_poll_for_query_results_exception(self):
        # Set expected values
        expected_query_execution_id = "expected-id"
        expected_exception = Exception("test")

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            "get_query_execution", QueryExecutionId=expected_query_execution_id
        ).raise_exception(expected_exception)

        # Test poll_for_query_results() exception
        a = Athena(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            a.poll_for_query_results(expected_query_execution_id)
        assert e_info.value == expected_exception
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_download_query_results_as_dict(self):
        # Set expected values
        expected_bucket = "bucket"
        expected_key = "key"
        expected_body = "test,csv,body\nalpha.val,,4.0\nbeta.val,-10,".encode("utf-8")

        # Set up mocks
        mock_s3_client = MockApi()
        mock_s3_client.on("get_object", Bucket=expected_bucket, Key=expected_key).return_val(
            {"Body": BytesIO(expected_body)}
        )
        mock_session, _ = self.get_mocks()
        mock_session.on("client", "s3").return_val(mock_s3_client)

        # Test download_query_results_as_dict()
        a = Athena(aws_session=mock_session)
        results = a.download_query_results_as_dict(
            {
                "QueryExecution": {
                    "ResultConfiguration": {
                        "OutputLocation": f"s3://{expected_bucket}/{expected_key}"
                    }
                }
            }
        )
        assert results == [
            {"test": "alpha.val", "csv": None, "body": "4.0"},
            {"test": "beta.val", "csv": "-10", "body": None},
        ]
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_download_query_results_as_dict_exception(self):
        # Set expected values
        expected_bucket = "bucket"
        expected_key = "key"
        expected_body = "test,csv,body\nalpha.val,-10,4.0,extra,columns".encode("utf-8")

        # Set up mocks
        mock_s3_client = MockApi()
        mock_s3_client.on("get_object", Bucket=expected_bucket, Key=expected_key).return_val(
            {"Body": BytesIO(expected_body)}
        )
        mock_session, _ = self.get_mocks()
        mock_session.on("client", "s3").return_val(mock_s3_client)

        # Test download_query_results_as_dict()
        a = Athena(aws_session=mock_session)
        with pytest.raises(AssertionError) as e_info:
            a.download_query_results_as_dict(
                {
                    "QueryExecution": {
                        "ResultConfiguration": {
                            "OutputLocation": f"s3://{expected_bucket}/{expected_key}"
                        }
                    }
                }
            )
        assert str(e_info.value).startswith("Invalid CSV column")
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_download_query_results(self):
        # Set expected values
        expected_bucket = "bucket"
        expected_key = "key"
        expected_body = "test,csv,body".encode("utf-8")

        # Set up mocks
        mock_s3_client = MockApi()
        mock_s3_client.on("get_object", Bucket=expected_bucket, Key=expected_key).return_val(
            {"Body": BytesIO(expected_body)}
        )
        mock_session, _ = self.get_mocks()
        mock_session.on("client", "s3").return_val(mock_s3_client)

        # Test download_query_results()
        a = Athena(aws_session=mock_session)
        results = a.download_query_results(
            {
                "QueryExecution": {
                    "ResultConfiguration": {
                        "OutputLocation": f"s3://{expected_bucket}/{expected_key}"
                    }
                }
            }
        )
        assert results == expected_body
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_download_query_results_exception(self):
        # Set expected values
        expected_bucket = "bucket"
        expected_key = "key"
        expected_exception = Exception("test exception")

        # Set up mocks
        mock_s3_client = MockApi()
        mock_s3_client.on("get_object", Bucket=expected_bucket, Key=expected_key).raise_exception(
            expected_exception
        )
        mock_session, _ = self.get_mocks()
        mock_session.on("client", "s3").return_val(mock_s3_client)

        # Test download_query_results() exception
        a = Athena(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            a.download_query_results(
                {
                    "QueryExecution": {
                        "ResultConfiguration": {
                            "OutputLocation": f"s3://{expected_bucket}/{expected_key}"
                        }
                    }
                }
            )
        assert e_info.value == expected_exception
        mock_s3_client.assert_expectations()
        mock_session.assert_expectations()

    @pytest.mark.unit
    def test_reload_partitions_exception(self):
        # Set expected values
        expected_table = "test_table"
        expected_exception = Exception("test")

        # Set up mocks
        mock_session, mock_athena_client = self.get_mocks()
        mock_athena_client.on(
            "start_query_execution",
            QueryString=f"MSCK REPAIR TABLE `{expected_table}`",
            QueryExecutionContext={"Database": "default", "Catalog": "AwsDataCatalog"},
            WorkGroup="primary",
        ).raise_exception(expected_exception)

        # Test reload_partitions() exception
        a = Athena(aws_session=mock_session)
        with pytest.raises(Exception) as e_info:
            a.reload_partitions(expected_table)
        assert e_info.value == expected_exception
        mock_athena_client.assert_expectations()
        mock_session.assert_expectations()

    def get_mocks(self) -> Tuple[MockApi, MockApi]:
        mock_athena_client = MockApi()
        mock_session = MockApi()
        mock_session.on("client", "athena").return_val(mock_athena_client)
        return mock_session, mock_athena_client
