import os
import snowflake.connector
from snowflake.sqlalchemy import URL
from sqlalchemy import create_engine
from absl import logging
from getpass import getpass

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


def get_snowflake_creds() -> dict:
    """
    Fetches and returns snowflake creds for connecting to snowflake

    Please use this within the scope of a function to ensure its memory safe 

    returns:
    - creds (dict) - a dictionary containing user creds

    """
    username = os.environ.get("SNOWFLAKE_USER", "aadamu")
    account = os.environ.get("SNOWFLAKE_ACCOUNT", "orchard")
    warehouse = os.environ.get("SNOWFLAKE_WAREHOUSE", "DEV_PERFORMANCE_WAREHOUSE")
    snowflake_password = os.environ.get("SNOWFLAKE_PASSWORD", getpass("\n>>Enter password:"))
    snowflake_passcode = getpass("\n>>Enter MFA Passcode:") or None
    creds = {
        "user":  username,
        "password": snowflake_password,
        "account": account,
        "warehouse": warehouse,
        "protocol": 'https',
        "passcode": snowflake_passcode
    }
    return creds


def alchemy_engine_factory(creds=None, **kwargs):
    try:
        if creds is None:
            _creds = get_snowflake_creds()
        else:
            _creds = creds
        # generate connection string
        conn_string = URL(
            account=_creds['account'],
            user=_creds['user'],
            password=_creds['password'],
            database=SNOWFLAKE_DB,
            schema=SNOWFLAKE_SCHEMA,
            warehouse=SNOWFLAKE_WAREHOUSE,
            role=SNOWFLAKE_ROLE
        )

        # init engine
        engine = create_engine(conn_string)
        return engine
    except Exception as e:
        logging.error(f"Something went wrong while creating engine - {str(e)}")


def alchemy_connector_factory(creds=None, **kwargs):
    """
     A Factory for creating SQL Alchemy Connectors

    This returns the cursor after opening a session with snowflake.

    params:
        - creds - snowflake credentials

    returns:
        - cursor - sql alchemy connection
    """

    try:
        engine = alchemy_engine_factory(creds, **kwargs)

        return engine.connect()
    except Exception as e:
        logging.error(f"Something went wrong - {str(e)}")


def snowflake_connector_factory(creds=None, return_as_cursor=True):
    """
    A Factory for creating snowflake connectors.

    This returns the cursor after opening a session with snowflake.

    params:
        - creds - snowflake credentials

    returns:
        - cursor - snowflake session cursor
    """
    try:
        if creds is None:
            _creds = get_snowflake_creds()
        else:
            _creds = creds
        # connect to snowflake
        conn = snowflake.connector.connect(**_creds)
        if return_as_cursor:
            return conn.cursor()
        else:
            return conn
    except Exception as e:
        logging.error(f"Something went wrong - {str(e)}")


def set_snowflake_environment(conn_cursor, warehouse="DEV_OWS_WAREHOUSE",
                database=None,
                schema=None):
    """ Setups Environment"""
    conn_cursor.execute(f"USE WAREHOUSE {warehouse};")
    if database:
        conn_cursor.execute(f"USE DATABASE {database};")
    if schema:
        conn_cursor.execute(f"USE SCHEMA {database}.{schema};")