"""
Pre Release Forecasting Pipleines

"""
from forecasting_toolkit.models.pipelines.inference.dataset.rollforward import (
    get_inference_dataset_by_rolling_forward_latest_track
)
from forecasting_toolkit.datastore.adapters.helpers import (
    export_to_snowflake_table,
    create_snowflake_table,
)
from forecasting_toolkit.datastore.connectors.snowflake import (
    snowflake_connector_factory,
    set_snowflake_environment
)
from snowflake.connector.pandas_tools import (
    write_pandas
)

from forecasting_toolkit.utils import (
    load_sklearn_model
)

import os
import argparse

import pandas as pd
from absl import logging
import mlflow 
from typing import List

# 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")

# set tracking uri
TRACKING_SERVER = os.environ.get("TRACK_URI", "https://dev-orch-mlflow-service.dev.theorchard.io/")
mlflow.set_tracking_uri(TRACKING_SERVER)


def inference_pipeline(store_id:int,
                        model_path,
                        country_code_id=1,
                        target_col:str = "STREAMS") -> List[pd.DataFrame]:
    """
    Inference pipeline runs predictions using the given model for the store

    Args:
        store_id (_type_): This is the store id from FACTS.PROD.DIM_STORE
        model_path (_type_): path to the model to load
        country_code_id (int, optional): Country to run the predictions for.
                                        Defaults to 1.
        target_col (str, optional): Target column. Defaults to "STREAMS".

    Returns:
        List[pd.DataFrame]: List of pandas dataframe.
                            (i.e. predictions for each track)
    """
    # load model
    logging.debug(f"Loading model at {model_path}")
    model = load_sklearn_model(model_path=model_path)
    logging.debug(model)

    forecasts_lst = []
    # For each track, on release batch find latest track release and rollforward for predictions
    for track_df in get_inference_dataset_by_rolling_forward_latest_track(store_id=store_id,
                                                                        country_code_id=country_code_id):
        try:
            # predict
            forecasted_streams = model.predict(track_df)
            # update dataframe
            track_df[target_col] = forecasted_streams
            forecasts_lst.append(track_df)
        except Exception as e:
            logging.warning(f'Data Issue while predicting - {str(e)}')
            continue

    return forecasts_lst


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description="Rolls forward track data" +
                                        "and runs model to generate forecasts")

    parser.add_argument('--store_id', type=int,
                            help='Store ID',
                            dest="store_id")

    parser.add_argument('--country_code_id', type=int,
                            help='Country Code ID',
                            default=1,
                            dest="country_code_id")

    parser.add_argument('--model_path',  type=str,
                            help='Model file to load',
                            dest="model_path",
                            required=True)

    parser.add_argument('--csv',  type=str,
                            default="./forecast_df.csv",
                            help='Forecast output',
                            dest="forecast_out",
                            required=False)

    parser.add_argument('--snowflake_table',  type=str,
                            default="DEBUT_FORECASTS",
                            dest="snowflake_output_table",
                            required=False)

    parser.add_argument('--debug', default=True,
                            help='Sets logging to debug mode',
                            dest="debug",
                            required=False)

    # parse arguments
    args = parser.parse_args()
    columns = [
            'SNAPSHOT_DATE', 'ISRC', 'TRACKNAME',
            'UPC', 'RELEASE_NAME', 'STORE_ID', 'FEED_ID',
            'RELEASE_DATE', 'STREAMS'
    ]
    if args.debug:
        logging.set_verbosity(logging.DEBUG)

    # run inference on the next batch of releases in STG_FORECAST
    logging.debug("Running inference with GBT Regressor V2")
    lst_forecasts = inference_pipeline(store_id=args.store_id,
                                                country_code_id=args.country_code_id,
                                                model_path=args.model_path)
    # combing forecasts into one dataframe
    logging.debug('Combing track forecasts')
    forecast_df = pd.concat(lst_forecasts, ignore_index=True)

    if args.forecast_out:
        logging.debug("Exporting predictions to csv file")
        # write out to file
        try:
            forecast_df[columns].to_csv(args.forecast_out, index=False)
        except Exception as err:
            logging.error(f"What a bummer - Error while writing to csv: {str(err)}")

    if args.snowflake_output_table:
        logging.debug("Exporting to snowflake")

        # clean up name
        snowflake_table = args.snowflake_output_table.strip()

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

        forecast_df.columns = [str(col).upper() for col in forecast_df.columns]
        forecast_df[columns].to_csv("pre_export_final.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=forecast_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)}")