"""
Snowflake adapter class definition

Table of Contents
- Imports
- Class definition
"""

#############
## Imports ##
#############

# Standard library imports
import pandas as pd
from typing import Any, Dict, Union
import psycopg2

# Prism-specific imports
from .adapter import Adapter
import prism.exceptions

# From https://docs.aws.amazon.com/redshift/latest/mgmt/python-configuration-options.html
REDSHIFT_PSYCOPG_CONFIG_OPTIONS = {
    'database': str,
    'host': str,
    'password': str,
    'port': int,
    'user': str
}

ADDITIONAL_CONFIGS = {
    'autocommit': bool
}

######################
## Class definition ##
######################

class Redshift(Adapter):

    def is_valid_config(self,
        config_dict: Dict[str, str]
    ) -> bool:
        """
        Check that config dictionary is profile.yml is valid

        args:
            config_dict: config dictionary under Redshift adapter in profile.yml
        returns:
            boolean indicating whether config dictionary in profile.yml is valid
        """

        # Required config vars
        required_config_vars = [
            'user',
            'password',
            'port',
            'host',
            'database'
        ]

        # Optional config vars
        optional_config_vars = [
            'autocommit'
        ]

        # Raise an error if:
        #   1. Config doesn't contain any of the required vars or contains additional config vars
        #   2. Any of the config values are None
        actual_config_vars = []
        for k,v in config_dict.items():
            if k not in required_config_vars and k not in optional_config_vars:
                raise prism.exceptions.InvalidProfileException(message=f'invalid var `{k}` under redshift config in profile.yml')
            if k in required_config_vars:
                actual_config_vars.append(k)
            if v is None:
                raise prism.exceptions.InvalidProfileException(message=f'var `{k}` under redshift config cannot be None in profile.yml')
        vars_not_defined = list(set(required_config_vars) - set(actual_config_vars))
        if len(vars_not_defined)>0:
            v = vars_not_defined.pop()
            raise prism.exceptions.InvalidProfileException(message=f'need to define `{v}` under redshift config in profile.yml')

        # If no exception has been raised, return True
        return True


    def create_engine(self,
        adapter_dict: Dict[str, Any],
        adapter_type: str
    ):
        """
        Parse Redshift adapter, represented as a dict and return the Redshift connector object

        args:
            adapter_dict: Redshift adapter represented as a dictionary
        returns:
            Redshift connector object
        """
        
        # Get configuration and check if config is valid
        self.is_valid_config(adapter_dict)

        # Create psycopg2 connection
        conn = psycopg2.connect(
            dbname=adapter_dict['database'],
            host=adapter_dict['host'],
            port=adapter_dict['port'],
            user=adapter_dict['user'],
            password=adapter_dict['password']
        )
        
        # Autocommit. If no autocommit is specified, then set to True
        try:
            autocommit_config = bool(adapter_dict['autocommit'])
            conn.set_session(autocommit=autocommit_config)
        except KeyError:
            conn.set_session(autocommit=True)
        return conn
        

    def execute_sql(self, query: str, return_type: str) -> pd.DataFrame:
        """
        Execute the SQL query
        """
        # Create cursor for every SQL query -- this ensures thread safety
        cursor = self.engine.cursor()
        cursor.execute(query)
        if return_type=="pandas":
            data = cursor.fetchall()
            cols = []
            for elts in cursor.description:
                cols.append(elts[0])
            df: pd.DataFrame = pd.DataFrame(data=data, columns=cols)
            cursor.close()
            return df
        else:
            cursor.close
        

# EOF