"""
Snowflake Warehouse Python connector context manager
====================================================

Context manager for connecting to Snowflake and helpful wrappers

Usage:
 with connect(user='{user}',password='{password}',account='{account}') as conn:
    with cursor(conn) as curs:
       curs.execute("SELECT current_version()")
       one = curs.fetchone()
       print(one[0])
"""

import base64
import contextlib
import os

from snowflake import connector


TABLE_EXISTS_SQL = "SHOW TABLES LIKE '{table}' in {db}.{schema}"


class FetchEnum(object):
    ALL = 1
    ONE = 2
    MANY = 3


DEFAULT_DB_CONFIG = {
    'user': os.environ.get('SNOWFLAKE_USER'),
    'password': os.environ.get('SNOWFLAKE_PASSWORD'),
    'account': os.environ.get('SNOWFLAKE_ACCOUNT'),
    'role': os.environ.get('SNOWFLAKE_ROLE'),
    'warehouse': os.environ.get('SNOWFLAKE_WAREHOUSE'),
    'db': os.environ.get('SNOWFLAKE_DATABASE'),
    'schema': os.environ.get('SNOWFLAKE_SCHEMA')
}

private_key = os.environ.get('SNOWFLAKE_KEY')
if private_key:
    decoded_key = base64.b64decode(bytes(private_key, encoding='utf-8'))
    DEFAULT_DB_CONFIG['private_key'] = decoded_key


@contextlib.contextmanager
def connect(**credentials):
    """Connection context manager for Snowflake.

    Returns:
        snowflake.SnowflakeConnection: connection object
    """

    ctx = connector.connect(**credentials)
    try:
        yield ctx
    finally:
        ctx.close()


@contextlib.contextmanager
def cursor(connection):
    """Cursor context manager that ensures closure on exit.

    Returns:
        snowflake.SnowflakeCursor: cursor instance
    """

    cs = connection.cursor()
    try:
        yield cs
    finally:
        cs.close()


def execute_with_py_conn(
        sql_statement, fetch_action=FetchEnum.ONE, sf_config=DEFAULT_DB_CONFIG,
        result_key='results'):
    """Executes sql on Snowflake, using context manager.

    Args:
        sql_statement (str): SQL to be executed
        fetch_action (int): fetch action (fetch one, many, etc.)
        sf_config (dict) optional: snowflake configuration and credentials
        result_key (str) optional: a key to form result dictionary

    Returns:
        dict: cursor row results
    """

    res = None

    with connect(
            user=sf_config.get('user'),
            password=sf_config.get('password'),
            account=sf_config.get('account'),
            role=sf_config.get('role'),
            private_key=sf_config.get('private_key')) as conn:
        with cursor(conn) as curs:
            curs.execute('USE WAREHOUSE {};'.format(
                sf_config.get('warehouse')))
            curs.execute('USE DATABASE {};'.format(
                sf_config.get('db')))
            curs.execute('USE SCHEMA {};'.format(
                sf_config.get('schema')))

            curs.execute(sql_statement)

            if fetch_action == FetchEnum.ALL:
                res = curs.fetchall()
            elif fetch_action == FetchEnum.ONE:
                res = curs.fetchone()

            result = dict()
            result[result_key] = res
            return result


def table_exists(table, sf_config=DEFAULT_DB_CONFIG):
    """Checks if a particular table exists in Snowflake.

    Args:
        table (str): name of the snowflake table
        sf_config (dict) optional: snowflake configuration and credentials

    Returns:
        bool: table exists
    """

    sql_statement = TABLE_EXISTS_SQL.format(
        db=sf_config.get('db'),
        schema=sf_config.get('schema'), table=table)

    result_key = 'exists'
    res = execute_with_py_conn(
        sql_statement, FetchEnum.ONE, sf_config, result_key)

    return bool(res.get(result_key))