import asyncio
from unittest.mock import MagicMock, Mock

import pytest
from pandas import DataFrame

from monday_com_orca_backend import config
from monday_com_orca_backend.connectors.snowflake_db import client

SAMPLE_ROW_LIMIT: int = 50
SAMPLE_QUERY = (
    "SELECT LABELID, LABELNAME FROM FACTS.PROD.DIM_LABEL LIMIT %s",
    (SAMPLE_ROW_LIMIT,),
)


@pytest.fixture(autouse=True)
def mute_logger(mocker):
    """Mute the logger for the Snowflake client."""
    mocker.patch.object(client, "logger")


class TestSnowflake:
    _class = client.Client
    limit_rows = 50

    @pytest.fixture
    def instance(self):
        """Create a Snowflake instance with real credentials,
        using the FACTS database and PROD schema for testing.
        """
        instance = self._class(**config.SNOWFLAKE_CREDENTIALS)
        instance.echo_interval = 0
        yield instance
        instance.disconnect()

    class TestAfetchOne:
        @pytest.mark.asyncio
        async def sample_query(self, instance):
            return await instance.afetch_one(*SAMPLE_QUERY)

        @pytest.fixture
        def mock_cursor_token_expired(self):
            return MagicMock(
                get_results_from_sfqid=Mock(
                    side_effect=client.sf_connector.errors.ProgrammingError(
                        "Token expired"
                    )
                ),
            )

        def assert_result(self, result):
            """Helper to assert the result is a dictionary with expected keys."""
            assert isinstance(
                result, dict
            ), "Result should be a dictionary (use DictCursor)"
            assert len(result) == 2, "Expected 2 columns in the result"

        @pytest.mark.asyncio
        async def test_afetch_one_multiple_sequential(self, instance):
            # Execute 2 to cover connection reuse
            for _ in range(2):
                result = await self.sample_query(instance)
                self.assert_result(result)

        @pytest.mark.asyncio
        async def test_afetch_one_multiple_concurrent(self, instance):
            tasks = [self.sample_query(instance) for _ in range(3)]
            results = await asyncio.gather(*tasks)
            for result in results:
                self.assert_result(result)

        @pytest.mark.asyncio
        async def test_afetch_one_with_expired_token_do_silent_reauth_and_retry(
            self, mocker, instance, mock_cursor_token_expired
        ):
            """Test silent reauthentication and retry when the token has expired."""
            mocker.patch.object(
                instance,
                "connection",
                cursor=Mock(return_value=mock_cursor_token_expired),
            )
            result = await self.sample_query(instance)
            self.assert_result(result)

        @pytest.mark.asyncio
        async def test_afetch_one_with_expired_token_prevent_infinite_recursivity(
            self, mocker, instance, mock_cursor_token_expired
        ):
            """Test silent reauthentication and retry when the token has expired."""
            mocker.patch.object(
                instance, "cursor", return_value=mock_cursor_token_expired
            )
            with pytest.raises(client.sf_connector.errors.ProgrammingError):
                await self.sample_query(instance)

    @pytest.mark.asyncio
    @pytest.mark.parametrize(
        "as_df,expected_result_type",
        [
            (True, DataFrame),
            (False, list),
        ],
    )
    async def test_afetch_all(self, instance, as_df, expected_result_type):
        result = await instance.afetch_all(*SAMPLE_QUERY, as_df=as_df)
        assert type(result) is expected_result_type
        assert len(result) == SAMPLE_ROW_LIMIT

    @pytest.mark.asyncio
    async def test_may_pass_database_and_schema_when_instantiating(self):
        instance = self._class(
            **config.SNOWFLAKE_CREDENTIALS,
            database="FACTS",
            schema="PROD",
        )
        # Because we passed the database and schema, we can use a query
        # without a fully qualified table name, as long as the table exists
        # in the specified database and schema.
        sample_query = (
            "SELECT LABELID, LABELNAME FROM DIM_LABEL LIMIT %s",
            (SAMPLE_ROW_LIMIT,),
        )
        result = await instance.afetch_all(*sample_query, as_df=False)
        assert len(result) == SAMPLE_ROW_LIMIT