import psycopg2
from contextlib import closing
import sys
import json

database = ""
user = ""
password = ""
host = ""
port = 5432

sql_query = {
    'unit_of_work': "select * from unit_of_work where activity_status='IN_PROGRESS' and now() - last_updated_at > interval '4 hours' order by last_updated_at;",
    'apps_etl_unit_of_work': "select * from apps_etl_unit_of_work where activity_status='IN_PROGRESS' and now() - last_updated_at > interval '4 hours' and unit_of_work_type<>'BACKFILL' order by last_updated_at;"
}




# grab sql data with closing connection after it
def fetch_sql_data(sql_req):
    with closing(psycopg2.connect(database=database, user=user, password=password, host=host, port=port)) as connection:
        with connection.cursor() as cursor:
            cursor.execute(sql_req)
            columns = cursor.description
            rows = cursor.fetchall()
            return columns, rows


# prepare list of dict arrays from sql result
def sql_to_list_dict(columns, rows) -> list:
    result = []
    for row in rows:
        tmp = {}
        for i in range(len(columns)):
            tmp[columns[i][0]] = row[i]
        result.append(tmp)
    return result


###################################################################################################
################################## filters functions ##############################################
# Just for info. Below functions works with list objects from outer space,
# and used to remove filtered information from eventual result message
# that sends to zabbix agent as a script print output
###################################################################################################

################ apps_etl_unit_of_work filters ################
###############################################################
# filter data according to active time difference and uow_code
def check_active_time(l_dict):
    l_dict_to_remove_list = []
    uow_code_list = ['youtubereporting']    # list of uow_codes that have to be filtered
    for row in l_dict:
        hours_difference = int((row['last_updated_at'].timestamp() - row['created_at'].timestamp()) / 3600)
        uow_code = row['unit_of_work_code'].split('-')[0]
        if hours_difference <= 15 and uow_code in uow_code_list:
            l_dict_to_remove_list.append(row)
        else:
            pass
    # if l_dict_to_remove_list isn't empty remove valid elements from list
    if l_dict_to_remove_list:
        for row in l_dict_to_remove_list:
            l_dict.remove(row)
###############################################################
###############################################################

################ unit_of_work #################################
###############################################################
# no filters here yet
###############################################################
###############################################################

###################################################################################################
###################################################################################################
###################################################################################################


if __name__ == '__main__':
    if len(sys.argv) > 1:
        arg = str(sys.argv[1])
        columns, rows = fetch_sql_data(sql_query[arg])
        data_list = sql_to_list_dict(columns, rows)
        ################ perform some check filters here according to sql query type ################
        if arg == 'unit_of_work':
            pass    # no filters here, since we don't have such requirements
        if arg == 'apps_etl_unit_of_work':
            check_active_time(data_list)
        ###########################################################################################

        ################ print result ################
        if len(data_list) <= 150:
            for row in data_list:
                tmp_dict = {key: str(value) for (key, value) in row.items()}    # make all value as str just for unboxing 'datetime.datetime' objects
                print(json.dumps(tmp_dict, indent=1))
        else:
            print("query output is too big, better to check it via DB client")
    else:
        print("no arguments were given")