import re

import pandas as pd
import pytest
import pytest_asyncio

from src.backend.logic.look_data.io import snowflake

LookTask = snowflake.LookTask


class TestSnowflakeRequest:
    _class = snowflake.SnowflakeRequest
    test_look_name: str = "TEST"

    # Limit the number of rows to fetch for testing purposes, for increased speed.
    fetch_row_count: int = 10

    sample_query: str = f"""
    SELECT RELEASEID
    FROM FACTS.PROD.DIM_RELEASE
    LIMIT {fetch_row_count}
    """

    @pytest.fixture(scope="class")
    def instance(self):
        return self._class(formatters={"RELEASEID": lambda x: str(x)[0:3]})

    @pytest_asyncio.fixture(scope="class")
    async def execute(self, instance):
        result = await instance.execute(
            [LookTask(look_name=self.test_look_name, query=self.sample_query)]
        )
        df = result[0]
        assert (
            len(df) == self.fetch_row_count
        ), f"Expected {self.fetch_row_count} rows, got {len(df)}"
        return df

    def test_formatters(self, execute):
        df = execute
        assert all(re.fullmatch(r"\d{3}", value) for value in df["RELEASEID"])

    @pytest.mark.asyncio
    async def test_execute(self, execute):
        df = execute
        assert isinstance(df, pd.DataFrame)
        assert len(df) == self.fetch_row_count

    @pytest.mark.asyncio
    async def test_execute_preprocessor(self, instance):
        original_column_name: str = "RELEASEID"
        expected_column_name: str = "CHANGED_RELEASEID"

        def preprocessor(df_: pd.DataFrame) -> pd.DataFrame:
            return df_.rename(columns={original_column_name: expected_column_name})

        result = await instance.execute(
            [
                LookTask(
                    look_name=self.test_look_name,
                    query=self.sample_query,
                    preprocessor=preprocessor,
                )
            ]
        )

        df = result[0]
        assert len(df) == self.fetch_row_count
        assert original_column_name not in df.columns
        assert expected_column_name in df.columns