#! /usr/bin/env python3

from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
import smtplib
import sql_queries
import postgres_reporter
import snowflake_reporter
import time
import argparse
import sys


def create_arg_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', '--mail',
                        nargs='+',
                        action='store',
                        dest='recipients',
                        type=str)
    parser.add_argument('-t', '--type_check',
                        nargs='+',
                        action='store',
                        dest='type_check',
                        required=True,
                        type=str)
    parser.add_argument('-p', '--print', action="store_const", const=True)
    return parser


def send_mail(recipients: str, message: str, sbj_type: str) -> None:
    # Create a multipart message
    msg = MIMEMultipart()
    body_part = MIMEText(message, 'plain')
    msg['Subject'] = f"[{sbj_type.upper()}] Daily Data Transfer status - Data Consistency"
    msg['From'] = "support.sme@dataart.com"
    msg['To'] = recipients
    # Add body to email
    msg.attach(body_part)
    # Create SMTP object
    server = smtplib.SMTP('relay1.dataart.com', 25)
    server.sendmail(msg['From'], msg['To'].split(','), msg.as_string())
    server.quit()


def decibel_ads_check():
    result_data = {}
    check_name_list = [x for x in sql_queries.postgres_ads.keys()]
    for key in check_name_list:
        time.sleep(2)
        # get the data from databases
        pg_data = postgres_reporter.fetch_sql_data(sql_queries.postgres_ads[key])
        sf_data = snowflake_reporter.snowflake_fetch_data(sql_queries.snowflake_ads[key], sf_database="DELPHI_ADS_DATA")
        # prepare variables for proper condition
        sf_data = int(sf_data[0]['COUNT(*)'])
        pg_data = int(pg_data[0]['count'])
        #
        if sf_data != pg_data:
            result_data[key] = {'snowflake': sf_data, 'postgres': pg_data}
    return result_data


def decibel_linkfire_check():
    result_data = {}
    # check the step with LINKFIRE_IMPORT_LOG
    linkfire_import_log = snowflake_reporter.snowflake_fetch_data(sql_queries.snowflake_linkfire["LINKFIRE_IMPORT_LOG"], sf_database="DELPHI_LINKFIRE")
    if int(linkfire_import_log[0]['COUNT(*)']) < 1:
        result_data["LINKFIRE_IMPORT_LOG"] = linkfire_import_log[0]['COUNT(*)']
    #
    # check the rest queries counts between snowflake and postgres
    check_name_list = [x for x in sql_queries.postgres_linkfire.keys()]
    for key in check_name_list:
        time.sleep(2)
        # get the data from databases
        pg_data = postgres_reporter.fetch_sql_data(sql_queries.postgres_linkfire[key])
        sf_data = snowflake_reporter.snowflake_fetch_data(sql_queries.snowflake_linkfire[key], sf_database="DELPHI_LINKFIRE")
        #
        # prepare variables for proper condition
        sf_data = int(sf_data[0]['COUNT'])
        if pg_data[0].get('count'):
            pg_data = int(pg_data[0]['count'])
        elif pg_data[0].get('sum'):
            pg_data = int(pg_data[0]['sum'])
        #
        if sf_data != pg_data:
            result_data[key] = {'snowflake': sf_data, 'postgres': pg_data}
    return result_data


if __name__ == '__main__':
    def prepare_message(data_dict):
        message = "COUNTS mismatching in the next check_names:\n\n"
        for key in data_dict:
            message += f"\n{key}: {data_dict[key]}\n"
        return message

    # # check command line parameters and print or email result
    parser = create_arg_parser()
    arg_list = parser.parse_args()
    message_final = ""
    type_check = ""
    if len(sys.argv[1:]) > 0:
        if arg_list.type_check:
            type_check = arg_list.type_check[0]
            if type_check == "ads":
                result_dict = decibel_ads_check()
                message_final = prepare_message(result_dict)
            elif type_check == "linkfire":
                result_dict = decibel_linkfire_check()
                message_final = prepare_message(result_dict)
        if arg_list.recipients:
            recipients = arg_list.recipients[0]
            send_mail(recipients=recipients, message=message_final, sbj_type=type_check)
        if arg_list.print:
            print(f"{type_check.upper()}\n{message_final}")
    else:
        parser.print_help()