import asyncio
from dataclasses import dataclass
from time import time
from typing import Callable

from pandas import DataFrame

from .... import logger
from ....config import SNOWFLAKE_CREDENTIALS
from ....connectors import snowflake_db

logger = logger.new_logger(__name__)


@dataclass(slots=True)
class LookTask:
    """Dataclass to hold a Look query and its
    preprocessor function, if any.

    Attributes:
        look_name: Name of the Look for which the query is the underlying
            SQL query. This is used for logging and debugging purposes.
            (e.g. "SR1", "SR2", etc.)
        query: Look underlying SQL query.
        preprocessor: Hook to preprocess the data from the query.

    """

    look_name: str
    query: str
    preprocessor: Callable = None


class SnowflakeRequest:
    # Limit the number of concurrent queries to prevent overloading Snowflake.
    max_concurrent_queries: int = 5

    # Snowflake credentials.
    credentials: dict[str, str] = {
        **SNOWFLAKE_CREDENTIALS,
        "database": "ORCHARD_APP_REPORTING_V2",
        "schema": "ART_RELATIONS_PROD_ART_RELATIONS",
    }

    def __init__(self, formatters: dict[str, Callable] = None):
        self._snowflake_client = None
        self._formatters = formatters or None
        self._snowflake_semaphore = asyncio.Semaphore(self.max_concurrent_queries)

    async def execute(
        self,
        look_tasks: list[LookTask],
    ) -> list[DataFrame]:
        """Execute a list of Look tasks in parallel.

        Args:
            look_tasks: List of LookTask objects.

        Returns:
            List of dataframes, matching the order of the look_tasks list.
        """
        if self._snowflake_client is None:
            self._new_client()

        results = await asyncio.gather(*[self._new_worker(task) for task in look_tasks])
        return results  # noqa

    def _new_client(self):
        """Instantiate Snowflake."""
        self._snowflake_client = snowflake_db.Client(**self.credentials)

    async def _new_worker(self, look_task: LookTask) -> DataFrame:
        """Query worker, used to run queries in parallel and apply
        transformations.

        Args:
            look_task: LookTask object.
        """
        look_name = look_task.look_name

        logger.debug(f"Look {look_name}: Fetching data from Snowflake...")
        initial_time = time()

        async with self._snowflake_semaphore:
            df = await self._snowflake_client.afetch_all(look_task.query, as_df=True)

        logger.debug(f"Look {look_name}: Data fetched in {time() - initial_time:.2f}s.")

        # Preprocessors must be run before formatters (for tasks
        # such as grouping and aggregating).
        if preprocessor := look_task.preprocessor:
            logger.debug(f"Look {look_name}: Running preprocessor on data...")
            curr_time = time()
            df = preprocessor(df)
            logger.debug(
                f"Look {look_name}: Preprocessing completed in {time() - curr_time:.2f}s."
            )
        else:
            logger.debug(f"Look {look_name}: No preprocessor to run.")

        if self._formatters:
            logger.debug(f"Look {look_name}: Applying formatters...")
            curr_time = time()
            self._apply_formatters(df)
            logger.debug(
                f"Look {look_name}: Formatting completed in {time() - curr_time:.2f}s."
            )
        else:
            logger.debug(f"Look {look_name}: No formatters to apply.")

        logger.debug(
            f"Look {look_name}: Worker completed job in {time() - initial_time:.2f}s."
        )
        return df

    def _apply_formatters(
        self,
        df: DataFrame,
    ) -> None:
        """Apply formatters in-place to the dataframe. Formatters are applied
        to the columns specified in the formatters dict, if they exist
        in the dataframe.
        """
        formatters = self._formatters
        columns_to_format = set(formatters.keys()).intersection(df.columns)
        for col in columns_to_format:
            df[col] = df[col].apply(formatters[col])