import os
import pandas as pd
from absl import logging
from typing import List
from forecasting_toolkit.datastore.adapters.snowflake import (
    SnowflakeDatasetAdapter
)
from forecasting_toolkit.datastore.connectors.snowflake import (
    snowflake_connector_factory,
    alchemy_engine_factory,
    set_snowflake_environment
)

from sqlalchemy import (
    MetaData,
    Table,
    Column,
    Integer,
    String,
    Float,
    Date
)

from snowflake.connector.pandas_tools import (
    write_pandas
)


# snowflake
SNOWFLAKE_WAREHOUSE = os.environ.get("SNOWFLAKE_WAREHOUSE", "DEV_OWS_WAREHOUSE")
SNOWFLAKE_DB = os.environ.get("SNWOFLAKE_DB", "DEV_ENGINEERING")
SNOWFLAKE_SCHEMA = os.environ.get("SNOWFLAKE_SCHEMA", "AADAMU_DEBUT_FORECASTING_DBT")
SNOWFLAKE_ROLE = os.environ.get("SNOWFLAKE_ROLE", "DEV_ENGINEERING")

""" Helpers """


def fetch_dataset(snowflake_table, filters={}, warehouse="DEV_OWS_WAREHOUSE"):
    # init dataset - empty dataset
    logging.debug("init empty dataset")
    dataset_df = pd.DataFrame([])

    # connect to snowflake
    logging.debug('Connecting to snowflake')

    with snowflake_connector_factory() as conn:
        # set snowflake environment
        logging.debug("setting up DB Env")
        set_snowflake_environment(conn_cursor=conn,
                warehouse=warehouse,
                database=SNOWFLAKE_DB,
                schema = SNOWFLAKE_SCHEMA)

        # create snowflake dataset adapter
        snowflake_dataset_adapter = SnowflakeDatasetAdapter(conn=conn)
        # fetch dataset 
        logging.debug("Fetching dataset from snowflake")
        dataset_df = snowflake_dataset_adapter.fetch_dataset(snowflake_table=snowflake_table,
                                                                filters=filters)
    return dataset_df


def fetch_dataset_with_connector(snowflake_table,
            snowflake_connector,
            filters={}):
    """
    Fetches and returns dataset using connector
    """
    # create snowflake dataset adapter
    snowflake_dataset_adapter = SnowflakeDatasetAdapter(conn=snowflake_connector)
    # fetch dataset 
    logging.debug("Fetching dataset from snowflake")
    dataset_df = snowflake_dataset_adapter.fetch_dataset(snowflake_table=snowflake_table,
                                                            filters=filters)
    return dataset_df


def create_snowflake_table(snowflake_table: str, target_col: str = 'STREAMS'):
    """
    Creates a table on snowflake
    columns = [
            'SNAPSHOT_DATE', 'ISRC', 'TRACKNAME',
            'UPC', 'STORE_ID', 'FEED_ID', 'RELEASE_DATE',
            'STREAMS'
        ]
    params:
        snowflake_table (str): snowflake table to create
        target_col (dict): snowflake table to drop
    """
    metadata_obj = MetaData()

    # Table - TODO: at a later point move this to a models module
    # TODO: At the moment snowflake has issues with write_pandas parsing
    # dates so we will use strings
    forecast_table = Table(
        snowflake_table,
        metadata_obj,
        Column('SNAPSHOT_DATE', String(90), nullable=False),
        Column('ISRC', String(90), nullable=False),
        Column('TRACKNAME', String(3000), nullable=True),
        Column('UPC', Integer, nullable=False),
        # Column('LABEL_ID', Integer, nullable=False),
        Column('STORE_ID', Integer, nullable=False),
        Column('FEED_ID', Integer, nullable=False),
        # Column('ARTIST_ID', Integer, nullable=False),
        # Column('ARTIST_NAME', String(90), nullable=True),
        Column('RELEASE_DATE', String(90), nullable=False),
        # Column('RELEASE_ID', String(90), nullable=False),
        Column('RELEASE_NAME', String(3000), nullable=True),
        Column('STREAMS', Float, nullable=True)
    )

    engine = alchemy_engine_factory()
    forecast_table.create(engine, checkfirst=True)


def export_to_snowflake_table(snowflake_table: str,
                            data_df: pd.DataFrame,
                            columns: List[str] = None,
                            target_col: str = 'STREAMS'):
    """
    Exports pandas dataframe to a snowflake tabl

    Args:
        snowflake_table (str): _description_
        data_df (pd.DataFrame): _description_
    """
    if columns is None:
        columns = [
            'SNAPSHOT_DATE', 'ISRC', 'TRACKNAME',
            'UPC', 'STORE_ID', 'FEED_ID', 'RELEASE_DATE',
            target_col
        ]
    # clean up name
    snowflake_table = snowflake_table.strip()

    # create snowflake table
    create_snowflake_table(snowflake_table=snowflake_table)
    data_df['SNAPSHOT_DATE'] = data_df['SNAPSHOT_DATE'].astype(str)
    data_df['RELEASE_DATE'] = data_df['RELEASE_DATE'].astype(str)

    data_df.columns = [str(col).upper() for col in data_df.columns]
    data_df.to_csv("pre_export.csv")
    with snowflake_connector_factory(return_as_cursor=False) as conn:
        try:
            # setup snowflake environment
            set_snowflake_environment(conn_cursor=conn.cursor(),
                                warehouse=SNOWFLAKE_WAREHOUSE,
                                database=SNOWFLAKE_DB,
                                schema=SNOWFLAKE_SCHEMA)
            # write to snowflake
            write_pandas(conn=conn,
                            df=data_df[columns],
                            table_name=str(snowflake_table).strip(),
                            database=SNOWFLAKE_DB,
                            schema=SNOWFLAKE_SCHEMA)

            logging.debug("Completed Writing to snowflake table!")
        except Exception as err:
            logging.error(f"Issue encountered while writing to snowflake: {str(err)}")