import pytest
from firebolt.async_db.connection import Connection, connect
from firebolt.async_db.cursor import Cursor
from prefect.testing.utilities import AsyncMock

import prefect_firebolt.database
from prefect_firebolt import FireboltDatabase
from prefect_firebolt.credentials import FireboltCredentials
from prefect_firebolt.database import query_firebolt


class TestFireboltDatabase:
    @pytest.mark.parametrize(
        "engine_name,engine_url",
        [
            ("thomas", None),
            (None, "http://firebolt.io/thomas"),
            ("thomas", "http://firebolt.io/thomas"),
        ],
    )
    def test_valid_engine_parameters(self, engine_name, engine_url):
        if engine_name and engine_url:
            with pytest.raises(ValueError):
                FireboltDatabase(
                    database="prod",
                    credentials=FireboltCredentials(token="abc123"),
                    engine_name=engine_name,
                    engine_url=engine_url,
                )
        else:
            FireboltDatabase(
                database="prod",
                credentials=FireboltCredentials(token="abc123"),
                engine_name=engine_name,
                engine_url=engine_url,
            )

    @pytest.mark.parametrize(
        "credentials",
        [
            (FireboltCredentials(username="me@example.com", password="hunter2")),
            (FireboltCredentials(token="abc123")),
        ],
    )
    async def test_get_connection(self, credentials, monkeypatch):
        connect_mock = AsyncMock()
        monkeypatch.setattr(prefect_firebolt.credentials, "connect", connect_mock)
        await FireboltDatabase(
            database="prod",
            credentials=credentials,
        ).get_connection()
        connect_mock.assert_called_once()

        call_kwargs = connect_mock.call_args.kwargs
        auth = call_kwargs.pop("auth")

        if credentials.token:
            assert auth.token == credentials.token.get_secret_value()
        else:
            assert auth.username == credentials.username
            assert auth.password == credentials.password.get_secret_value()
        assert call_kwargs == {
            "database": "prod",
            "engine_name": None,
            "engine_url": None,
            "api_endpoint": "api.app.firebolt.io",
            "additional_parameters": {},
        }


class TestQueryFirebolt:
    @pytest.fixture
    async def mock_connect(self, monkeypatch):

        # returns a mocked cursor for testing
        firebolt_connect = AsyncMock(spec=connect, name="mock_connect")
        connection = AsyncMock(spec=Connection, name="mock_connection")
        cursor = AsyncMock(spec=Cursor, name="mock_cursor")

        # link all the mocks together appropriately
        firebolt_connect.return_value.__aenter__.return_value = connection
        connection.cursor.return_value = cursor
        cursor.execute = AsyncMock()
        monkeypatch.setattr(prefect_firebolt.credentials, "connect", firebolt_connect)

        return connection, cursor

    async def test_execute_query(self, mock_connect):
        mock_connection, mock_cursor = mock_connect

        mock_cursor.fetchall.return_value = ["TESTDB"]
        result = await query_firebolt.fn(
            database=FireboltDatabase(
                credentials=FireboltCredentials(token="abc123"), database="TESTDB"
            ),
            query="SHOW DATABASES",
        )
        assert result == ["TESTDB"]
        mock_cursor.execute.assert_called_with(query="SHOW DATABASES", parameters=None)

    async def test_execute_query_with_parameters(self, mock_connect):
        mock_connection, mock_cursor = mock_connect

        mock_cursor.fetchall.return_value = ["TESTDB"]
        result = await query_firebolt.fn(
            database=FireboltDatabase(
                credentials=FireboltCredentials(token="abc123"), database="TESTDB"
            ),
            query="SHOW DATABASES",
            parameters=("foo", "bar"),
        )
        assert result == ["TESTDB"]
        mock_cursor.execute.assert_called_with(
            query="SHOW DATABASES", parameters=("foo", "bar")
        )
