import os
import sys
import json

import snowflake.connector
import pandas as pd

def get_creds(env, flavor='SF'):
    '''
    Gets credentials from enviornment variables from .env
    '''
    credentials = {}
    if flavor == 'SF':
        if env == 'env':
            credentials['username'] = os.environ.get('SF_USER')
            credentials['password'] = os.environ.get('SF_PASS')
            credentials['role'] = os.environ.get('SF_ROLE')
            credentials['account'] = os.environ.get('SF_ACCOUNT')

        elif os.path.exists(env):
            with open(env) as f:
                credentials = json.dumps(f)

        else:
            print("Datebase not recognized.")
            credentials = -1
            
    elif flavor == 'S3':
        if env == 'env':
            credentials['KEY_ID'] = os.environ.get('AWS_ACCESS_KEY_ID')
            credentials['SECRET'] = os.environ.get('AWS_SECRET_ACCESS_KEY')

        elif os.path.exists(env):
            with open(env) as f:
                credentials = json.dumps(f)

        else:
            print("Datebase not recognized.")
            credentials = -1     
    
    return credentials

def _underscore(string):
    '''
    converts a space-delimited column header to a lowercase, underscored 
    column.
    removes characters that can ruin SQL commands such as hyphens (-) and 
    periods(.).
    '''
    return string.lower().rstrip() \
        .lstrip().replace(" ","_") \
        .replace('-','_') \
        .replace('.','_')

def _iterator(execution, columns, dtype):
    '''
    Iterates through a pandas dataframe.
    '''
    next = True

    while next:

        data = [row for row in execution.fetchmany()]
        if data:
            yield pd.DataFrame(data=data,
                               columns=columns,
                               dtype=dtype)
        else:
            next = False
            
def _without(d, key):
    '''
    returns a dict without a key.
    this is generally used to return a response from the cnx without a 'df' key.
    '''
    new_d = d.copy()
    new_d.pop(key)
    return new_d

class connect(object):
    """
    A connector object that communicates with databases
    """
    def __init__(self, warehouse=None, credentials='env'):
        """
        Args:
            credentials (dict) of db properties generated from get_creds
            --- or ---
            credentials (str) name of Database to be fed into get_creds
            warehouse (string) which warehouse to use, only for Snowflake connection.
        """
        credentials = get_creds(credentials)
                    
        try:
            cnx = snowflake.connector.connect(
                user = credentials['username'],
                password = credentials['password'],
                account= credentials['account'],
                role= credentials['role'],
                warehouse= warehouse,
                autocommit= True)
            
            self.connection = cnx
            self.warehouse = warehouse

        except snowflake.connector.Error as err:
            print(err)
    
    def close(self):
        self.connection.close()


    def q(self, query, resp='infer', chunksize=-1, dtype=None):
        """
        Updated
        Wrapper function to execute SQL commands, and return fetched responses.

        args ~ 
        query (String) containing SQL commands (required).
        MUTE (boolean) set to True to not fetch a response from the connector (default False).

        steps ~
        1. Creates a cursor function from the connection to register commands.
        2. Seaches String for command actions
        3. Executes command.
        4. Fetches output based on the the command (unless MUTE = True).
        5. Catches common errors.

        """
        # 1
        clc = self.connection.cursor()
        clc.arraysize = chunksize
        
        single_responses = ['PUT','INSERT','CREATE','DROP','COPY INTO','TRUNCATE','UPDATE']
        list_responses   = ['SHOW','DESC','SELECT']
        
        execution = clc.execute(query)
        columns = [col[0] for col in clc.description]
        
        # chunksize
        if resp == 'iterator':
            return _iterator(execution, columns, dtype)

        # no response...
        if resp == False:
            return clc.execute(query)
        
        # list response...
        elif any(x in query.upper() for x in list_responses) and \
             any(y in query.upper() for y in single_responses) == False:

            data = [row for row in execution.fetchall()]
            
            if data:
                first = data[0]
            else:
                first = None

            df = pd.DataFrame(data=data,
                              columns=columns,
                              dtype=dtype)
        # single response..
        else:
            first= execution.fetchone()[0]
            data = []
            df = pd.DataFrame()
             
        response = {'df' : df,
                    'first': first,
                    'code': 200,
                    'sfqid': clc.sfqid}
        
        if resp == 'df':
            return df
        elif resp == 'first':
            return first
        elif resp == 'list':
            return data
        else:
            return response
        
    def create_table(self, ref, tbl, fmt, debug=False, dtype=None, 
                     delim=False, comment='Auto-generated by Python Snowflake2 module.',):
        '''
        reads in a csv from local disk and creates a table with approrpiate column 
        names and sql dtypes.
        takes params of _input which is a csv, the date format, and the header 
        column to skip.
        set debug to True to print output instead of sending through SF client.
        '''

        def lazy_col(ref, file_format, dtype=dtype, delim=delim):
            '''
            Used in create_table() to generate a list of column names and dtypes 
            for building Snowflake tables. Uses Pandas to read a sample of a local 
            csv and converts columns and dytpes to SQL-notation by removing
            troublesome chars and using a dict for dtype conversions.

            Args ~ 
            _input (String) absolute path of file (can be s3) to create table framework.
                            dataframes are also acccepted.
            file_format (String) absolute snowflake path (db.sch.format_name).

            Steps ~     
            1.  checks the file format file for the table, uses it to choose 
                delimiters, compression, and file type.
            2. reads in a sample of a csv either from first 10k records (fast) or 
            3.  scans the column headers for datetime keywords, and converts those 
                columns to datetime
            4.  returns a comma-separated string of column_nameX SQLdtypeX,
                 column_nameY SQLdtypeY, etc...
            '''
            # dictionary for Numpy -> SQL dtypes (used in lazy_col)
            pd_2_sql_dtypes = {"object": "VARCHAR",
                               "str": "VARCHAR",
                               "int64": "NUMBER",
                               "float64": "FLOAT",
                               "bool": "BOOLEAN",
                               "datetime64[ns]": "DATE",
                               "<M8[ns]": "DATE"}

            # search column names for these words to convert to datetime
            datetime_keywords = ['date', 'day'] 

            if not delim:
                # 1 get metadeta for the format specified in file_format        
                if 'CSV' in file_format.upper():
                    delim = ','
                elif 'TSV' in file_format.upper():
                    delim = '\t'
                elif 'JSON' in file_format.upper():
                    print("TODO")

            # 2a  
            if isinstance(ref, pd.DataFrame) == False:
                # assumes compression is denoted in the filename with '.gz'
                compression='infer'
                if '.gz' in ref:
                    compression='gzip'

            if isinstance(ref, pd.DataFrame):
                df = ref
            
            elif 's3://' in ref:
                import s3
                df = s3.read_csv(ref, nrows=50000, sep=delim, 
                                 compression=compression, dtype=dtype)
            else:
                df = pd.read_csv(ref, nrows=50000, sep=delim, 
                                 compression=compression, dtype=dtype)
            if isinstance(df, pd.DataFrame):

                # 3
                dates = [col for col in df.columns 
                        if any(dateType in col.lower() 
                        for dateType in datetime_keywords)]

                if any(dates):
                    for date in dates:
                        try:
                            df[date] = pd.to_datetime(df[date])
                        except:
                            continue
                #4            
                return ','.join( [(_underscore(key) +" "+ pd_2_sql_dtypes[str(item)]) 
                                 for (key,item) in df.dtypes.iteritems()] )


        action = """CREATE OR REPLACE TABLE {TABLE} 
                    COMMENT = '{COMMENT}'
                    ({COLUMNS})
                    STAGE_FILE_FORMAT = (FORMAT_NAME ={FORMAT_NAME})""".format(
                        TABLE= tbl, FORMAT_NAME= fmt, COMMENT= comment,
                        COLUMNS= lazy_col(ref, file_format= fmt, 
                                          dtype=dtype, delim=delim))
        if debug==False:
            return _without(self.q(action), 'df')
        else:
            return action
        
    def s3_2_table(self, s3_path, tbl, fmt, debug=False, credentials='env'):
        '''
        Puts a file stored in an s3 location and copies it into a Snowflake table.
        Uses format_name to parse the file.

        globals ~ 
        s3_creds (json file) containing AWS 'KEY_ID' and 'SECRET'

        args ~
        s3_path (string) absolute path of file in s3 (s3://bucket/key/**/filename.ext).
        tbl (string) absolute path of table in Snowflake (db.sch.tbl)
        fmt (string) absolute path os format_name in snowflake (db.sch.format_name)
        debug (bool) set to True to print output. defaugit alts to False.
        '''
 
        credentials = get_creds(flavor='S3', env=credentials)

        action = """COPY INTO {TABLE} FROM '{BUCKET}'
                    CREDENTIALS = (AWS_KEY_ID='{KEY}' AWS_SECRET_KEY='{SECRET}')
                    FILE_FORMAT = (FORMAT_NAME = {FORMAT_NAME})""".format(
                        TABLE= tbl, BUCKET= s3_path, KEY= credentials['KEY_ID'],
                        SECRET= credentials['SECRET'], FORMAT_NAME= fmt)

        if debug==False:     
            return _without(self.q(action), 'df')
        else:
            return action

    def count_rows(self, tbl, where=dict(), wh='SCIENCE'):
        """
        Counts number of rows in table
        given filter parameters determinged in the unique_val dictionary
        """
        rows = -1
        i = 0

        if isinstance(where, dict):
        
            q = """
            SELECT COUNT(*)
            FROM {}
            """.format(table)
            
            
            for column, value in where.items():
                if i == 0:
                    if '__date__' in column:
                        value = "TO_DATE('{}')".format(value)
                    q += " where {} = '{}'".format(column, value)
                    i += 1
                else:
                    if '__date__' in column:
                        value = "TO_DATE('{}')".format(value)
                    q += " AND {} = '{}'".format(column, value)
                    i += 1
            try:           
                df = self.q(q, 'df')
                
                if isinstance(df,pd.DataFrame):
                    # if df['COUNT(*)'] != 0:
                    rows = df['COUNT(*)'].iloc[0]
                
                else:
                    return 0
                return rows
            except Exception as e:
                print(e)

        else:
            print("the where argument must be a dictionary!")
    
    def insert(self, into_table, from_table, new_columns=dict()):
        """
        Insert values into the "into" table
        from the "from" table with additional
        columns from new_columns
        """

        q = "INSERT INTO {} SELECT *".format(into_table)
        
        if isinstance(new_columns, dict):

            for column, value in new_columns.items():

                q += ", '{}' AS {}".format(value, column)
            
            q += " FROM {}".format(from_table)
            
            return _without(self.q(q), 'df')

        else:
            print("new_columns is not a dictionary")
            return

    def drop(self, table):
        '''
        Drops a table
        '''
        q = 'DROP TABLE {}'.format(table)
        return _without(self.q(q), 'df')

    def truncate(self, table):
        '''
        Truncates a table
        '''
        q = 'TRUNCATE TABLE {}'.format(table)
        return _without(self.q(q), 'df')