import asyncio
from unittest.mock import MagicMock, Mock

import pytest
from pandas import DataFrame

from src.backend.config import SNOWFLAKE_CREDENTIALS
from src.backend.connectors import snowflake_db

snowflake_db.logger = Mock()

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


class TestSnowflake:
    _class = snowflake_db.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(
            **SNOWFLAKE_CREDENTIALS,
            database="FACTS",
            schema="PROD",
        )
        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=snowflake_db.client.sf_connector.errors.ProgrammingError(
                        "Token expired"
                    )
                ),
            )

        @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)
                assert isinstance(result, tuple)
                assert len(result) == 2

        @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:
                assert isinstance(result, tuple)
                assert len(result) == 2

        @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)
            assert isinstance(result, tuple)
            assert len(result) == 2

        @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(
                snowflake_db.client.sf_connector.errors.ProgrammingError
            ):
                await self.sample_query(instance)

    @pytest.mark.asyncio
    @pytest.mark.parametrize("as_df", [True, False])
    async def test_afetch_all(self, instance, as_df):
        result = await instance.afetch_all(*SAMPLE_QUERY, as_df=as_df)
        assert isinstance(result, DataFrame if as_df else list)
        assert len(result) == 50