"""
Snowflake Dataset Adapter

"""
from forecasting_toolkit.datastore.adapters.base import (
    BaseDatasetAdapter
)

from snowflake.connector.pandas_tools import pd_writer
import pandas as pd
from absl import logging


# Snowflake Dataset Adapter


class SnowflakeDatasetAdapter(BaseDatasetAdapter):

    def __init__(self, conn):
        self.connector = conn

    def _execute_sql(self, sql: str, limit: int = None):
        """ Executes SQL

        Args:
            sql (str): Executes sql statement
            limit (int, optional): limt the results fetched. Defaults to None.

        Returns:
            _type_: results and rows object
        """
        try:
            res = self.connector.execute(sql)
            if limit is not None:
                rows = res.fetchmany(limit)
            else:
                rows = res.fetchall()
            return res, rows
        except Exception as e:
            logging.debug(f"sql: {sql}")
            logging.error(f"Opps...something went wrong. {e}")

    def execute_sql(self, sql: str, limit: int = None):
        """ Executes SQL and returns data as pandas dataframe """
        try:
            res, rows = self._execute_sql(sql=sql,
                                            limit=limit)
            return pd.DataFrame(rows,
                                columns=map(lambda meta: meta[0], res.description))
        except Exception as e:
            logging.error(f"Opps...something went wrong while fetching data from snowflake : {e}")

    def fetch_dataset(self, snowflake_table: str,
                    filters: dict = {},
                    limit: int = None):
        """ 
        Fetches dataset

        params:
            - snowflake_table (str) - snowflake table
            - filters (dict) - filters

        returns:
            dataset_df (pd.DataFrame) - dataset
        """
        sql_str = self.generate_sql(snowflake_table=snowflake_table,
                                    filters=filters, limit=limit)
        # final sql statement
        logging.debug(f"final sql: {sql_str}")
        dataset_df = self.execute_sql(sql=sql_str,
                                        limit=None)
        return dataset_df

# Alechmy Dataset Adapter


class AlchemyDatasetAdapter(BaseDatasetAdapter):
    """
    Alchemy Dataset Adapter
    Uses SQL Alchemy which allows us to write pandas
    dataframes to snowflake
    """
    def __init__(self, conn):
        self.connector = conn

    def execute_sql(self, sql: str, limit:int = None):
        """ Executes SQL and returns data as pandas dataframe """
        try:
            res = pd.read_sql_query(sql=sql,
                        con=self.connector,
                        chunksize=limit)
            return pd.DataFrame(res)
        except Exception as e:
            logging.error(f"Hmm..Something went wrong while reading table from snowflake : {e}")
            logging.debug(sql)

    def fetch_dataset(self, snowflake_table: str,
                        filters: dict = {},
                        limit: int = None):
        """
        Fetches dataset

        params:
            - snowflake_table (str) - snowflake table
            - filters (dict) - filters

        returns:
            dataset_df (pd.DataFrame) - dataset
        """
        sql_str = self.generate_sql(snowflake_table=snowflake_table,
                                    filters=filters, limit=limit)
        # final sql statement
        logging.debug(f"final sql: {sql_str}")
        dataset_df = self.execute_sql(sql=sql_str,
                                        limit=None)
        return dataset_df

    def to_snowflake(self, snowflake_table: str,
                        data_df: pd.DataFrame,
                        if_exists: str = 'replace'):
        """
        Writes pandas dataframe to snowflake table

        params:
            snowflake_table (str): snowflake table to write to
            data (pd.DataFrame): pandas dataframe to write to snowflake table
            replace_ifexists (bool, optional): Whether to replace table
                if it already exists. Defaults to False.

        returns:
            None
        """
        logging.debug(f"Writting to snowflake table: {snowflake_table}")
        try:
            # write to snowflake table
            data_df.to_sql(name=snowflake_table.lower(),
                            con=self.connector,
                            if_exists=if_exists,
                            method=pd_writer,
                            index=False)

            logging.debug("Completed write")
        except Exception as err:
            logging.debug("Something went wrong while writting to snowflake")
            logging.error(str(err))