""" This module defines all exergi functions within the AWS.postgreSQL module"""

def config(filePath: str ='/home/ec2-user/SageMaker/.config/database.ini', 
    section: str = 'maximo')-> dict:
    """ The following config() function reads in the database.ini file and 
    returns the connection parameters as a dictionary. This function will be 
    imported in to the main python script. 
    
    This file was adapted from:
    http://www.postgresqltutorial.com/postgresql-python/connect/
    
    Arguments:
        - filePath      - Path to the "database.ini"-file required to initiate
                          connection.
        - section       - Section of "database.ini"-filefile where the 
                          connection-parameters are stored.
    Returns:
        - db_cred  - All connection parameters in the specified filePath   
                          and section
    """

    from configparser import ConfigParser

    parser = ConfigParser()     # Create a parser
    parser.read(filePath)       # Read config file
    db_cred = {}                # Get section, default to postgresql
    
    # Checks to see if section (postgresql) parser exists
    if parser.has_section(section):
        params = parser.items(section)
        for param in params:
            db_cred[param[0]] = param[1]
    
    # Returns an error if a parameter is not listed in the initialization file
    else:
        raise Exception(f'Section {section} not found in the {filePath} file')
    return db_cred

def checkPostgreSQLConnection(db_cred: dict):
    """A function that checks connection settings and version of postgreSQL
    
    Arguments:
        - db_cred        - Connection parameters from config file
    """

    import psycopg2

    try:
        connection = psycopg2.connect(**db_cred)
        cursor = connection.cursor()
        
        # Print PostgreSQL Connection properties
        print (connection.get_dsn_parameters(),"\n")  
        
        # Print PostgreSQL version
        cursor.execute("SELECT version();")           
        record = cursor.fetchone()
        print("You are connected to - ", record,"\n")

    except (Exception, psycopg2.Error) as error:
        print ("Error while connecting to PostgreSQL", error)
    finally:
        # Closing database connection.
            if(connection):
                cursor.close()
                connection.close()
                print("PostgreSQL connection is closed")
                
def importDataFromPostgreSQL(sql_query: str,db_cred: dict):
    """A function that takes in a PostgreSQL query and outputs a pandas database 
    
    Arguments:
        - sql_query     - SQL-query to run
        - db_cred        - Connection parameters from config file
    Returns:
        - df            - DataFrame with Loaded Data
    """

    import psycopg2
    import pandas as pd

    try:
        connection = psycopg2.connect(**db_cred)
        cursor = connection.cursor()
        df = pd.read_sql_query(sql_query, connection)
        
        # Convert pandas dtype="object" columns to pd.StringDtype()
        for col in df.select_dtypes("object").columns:
            df[col] = df[col].astype(pd.StringDtype())
            
    except (Exception, psycopg2.Error) as error:
        print ("Error while connecting to PostgreSQL", error)
    finally:
            if(connection):
                cursor.close()
                connection.close()
                print("PostgreSQL connection is closed")
    return df

def importData(sql_query: str, db_cred: dict, df_meta = None, 
               post_rename: dict = {}, convert_dtypes: bool = True, 
               remove_tz: bool = False):
    """
    NOTE: This function will replace importDataFromPostgreSQL in future releases
    
    A function that takes in a PostgreSQL query and outputs a pandas database 
    
    Arguments:
        - sql_query         - SQL-query to run
        - db_cred           - Connection parameters from config file
        - df_meta           - DataFrame with schema information from importMetaData function
        - post_rename       - Dictionary for renaming dictionaries after dtype conversion
        - convert_dtypes    - Flag if data types should be converted using the 
                              data base schema for the used table
        - remove_tz         - Flag if timezone info should be removed from datetime columns
    Returns:
        - df                - DataFrame with loaded data
    """

    import psycopg2
    import pandas as pd
    import numpy as np
    connection = None
    
    try:
        
        # Get Data
        connection = psycopg2.connect(**db_cred)
        cursor = connection.cursor()
        df = pd.read_sql_query(sql_query, connection)
        
        if (df_meta is not None) & (convert_dtypes):

            for col,data_dtype in df.dtypes.iteritems():
                
                # Check schema
                df_meta_subset = df_meta.loc[df_meta.column_name==col,]
                
                # Test if schema is compatible
                assert col in df_meta_subset.column_name.unique(), f"Column {col} not present in schema"
                #assert df_meta_subset.data_type.nunique() == 1, f"Column {col} has multiple dtypes in schema"  Uncomment once Alex has replied to email
                
                if df_meta_subset.data_type.nunique() > 1:
                    print(f"Warning! '{col}' has multiple dtypes in schema. Using most popular dtype for conversion")
                    dtype = df_meta_subset.data_type.value_counts().index[0]
                else:
                    dtype = df_meta_subset.data_type.unique().item()

                # Integer conversion
                if dtype == 'integer':
                    if data_dtype != np.dtype('int64'):
                        df[col] = df[col].astype(pd.Int64Dtype())    
                
                # String conversion
                elif dtype == 'character varying':
                    if data_dtype == np.dtype('O'):
                        df[col] = df[col].astype(pd.StringDtype()) 
                
                # DateTime conversion
                elif dtype == 'timestamp with time zone':
                    if data_dtype == np.dtype('O'):
                        df[col] = pd.to_datetime(df[col],infer_datetime_format=True,utc=True,errors="coerce")
                
                # Date conversion
                elif dtype == 'date':
                    if data_dtype == np.dtype('O'):
                        df[col] = pd.to_datetime(df[col],infer_datetime_format=True,utc=True,errors="coerce")

                # Floats conversions
                elif ((dtype == 'double precision') | (dtype == 'numeric')):
                    if data_dtype != np.dtype('float64'):
                        df[col] = pd.to_numeric(df[col], errors='coerce')
            
                # Boolean conversion
                elif dtype == 'boolean':
                    if data_dtype != np.dtype('bool'):
                        df[col] = df[col].astype(bool)
                
                else:
                    raise(Exception(f"Data conversion '{dtype}' not implemented. Please contact package administrator"))
        else:
            print("WARNING. No data type conversion was executed")
        
        # Remove Timezone info from datetime tags
        for col in df.select_dtypes(include='datetime64[ns, UTC]').columns:
            df[col] = df[col].dt.tz_localize(None)
        
    except (Exception, psycopg2.Error) as error:
        print ("Error while connecting to PostgreSQL", error)
    finally:
            if(connection):
                cursor.close()
                connection.close()
                print("PostgreSQL connection is closed")
    
    # Post Dtype renaming
    for col in post_rename:
        assert col in df.columns, f"Column {col} in renameing dict not present in DataFrame"
    df = df.rename(columns=post_rename) 
    
    # Final check of dtypes
    if convert_dtypes:
        incorr_dtypes = [col for col in df.columns if df[col].dtype == np.dtype('O')]
        assert (len(incorr_dtypes) == 0), f"Column(s) {','.join(incorr_dtypes)} are dtype Object"
    
    return df

def importMetaData(db_cred: dict,like_str: str = "maximoworker_"):
    """Function for importing data base schema for every table with a tablename 
    starting with like_str 
    
    (SQL: ... WHERE table_name LIKE '{like_str}%)
    
    Arguments:
    - like_str - List all tables with table_name starting with this string. 
                 (default: "maximoworker_")
    - db_cred  - Dictionary containing all login information either manually 
                 entered or using the config function of the exergi-package.
    Returns
    - df_meta - DataFrame containing 
    """
    
    import psycopg2
    import pandas as pd
    connection = None
    
    # Format Query
    df_meta_query = f"""
                SELECT table_name,column_name,data_type,is_nullable FROM 
                information_schema.columns WHERE table_name LIKE '{like_str}%'"""
    try:

        # Import Data
        connection = psycopg2.connect(**db_cred)
        cursor = connection.cursor()
        df_meta = pd.read_sql(df_meta_query,connection).drop_duplicates()

    except (Exception, psycopg2.Error) as error:
        
        print ("Error while connecting to PostgreSQL", error)
        
    finally:
            if(connection):
                cursor.close()
                connection.close()
                print("PostgreSQL connection is closed")
    try:
        # Check for inconsistencies 
        if df_meta.duplicated().sum() > 0:
            raise Exception("Data types not consistent in schema")
    except (Exception) as error:
        
        print ("Error while connecting to PostgreSQL", error)
        
    return df_meta