#! /usr/bin/env python3

from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from contextlib import closing
from sql import sql_queries
from datetime import date, timedelta
import snowflake.connector
import psycopg2
import smtplib
import argparse
import sys
import json
import datetime

with open('config.json', 'r') as f:
    config = json.load(f)
sf_user = config['snowflake']["user"]
sf_password = config['snowflake']["password"]
sf_account = config['snowflake']["account"]
sf_warehouse = config['snowflake']['warehouse']

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


def send_mail(recipients: str, message: str, status: str) -> None:
    # Create a multipart message
    msg = MIMEMultipart()
    body_part = MIMEText(message, 'plain')
    msg['Subject'] = f"[{status}][PROD] AE Retargeting Dashboard Monitoring"
    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 snowflake_fetch_data(sql_query) -> dict:
    ctx = snowflake.connector.connect(
        user=sf_user,
        password=sf_password,
        account=sf_account,
        warehouse=sf_warehouse
    )
    cs = ctx.cursor(snowflake.connector.DictCursor)
    try:
        cs.execute(sql_query)
        result_dict = cs.fetchall()
    finally:
        cs.close()
    ctx.close()
    return result_dict

def check_1_9() -> dict:
    result_dict = {}
    full_log = {}

    ### Check 1-4 ##################################################
    for query_name, query in sql_queries.sf_queries_1_4.items():
        tmp_check = snowflake_fetch_data(f"""{query}""")
        full_log[f"{query_name}"] = {}
        if tmp_check:
            expected = "SUCCESS"
            for item in tmp_check:
                result = item['result']
                print(f"{result}")
                if result != expected:
                    full_log[f"{query_name}"] = {"Failed: Not all results are SUCCESS"}
                    result_dict[f"{query_name}"] = full_log[f"{query_name}"]
                    break
                else:
                    full_log[f"{query_name}"] = {"OK: All results are SUCCESS"}
        else:
            full_log[f"{query_name}"] = {f"OK: Query returned nothing"}
##            result_dict[f"{query_name}"] = full_log[f"{query_name}"]

        ### Check 5-8 ##################################################
        for query_name, query in sql_queries.sf_queries_5_8.items():
            tmp_check = snowflake_fetch_data(f"""{query}""")
            full_log[f"{query_name}"] = {}
            if tmp_check:
                full_log[f"{query_name}"] = {"Failed: Query returned not empty result"}
                result_dict[f"{query_name}"] = full_log[f"{query_name}"]
            else:
                full_log[f"{query_name}"] = {f"OK: Query returned nothing"}

    return result_dict, full_log

################# run the all steps ######################
def run_steps() -> dict:
    message_dict = {'name':"[PROD] AE Retargeting Dashboard Monitoring",'status':"",'date':"", 'info':"", 'details':"", 'full_log':""}
    # set date according to sql searching date, currently it's `today`
    message_dict["date"] = message_dict["date"] = datetime.datetime.now().strftime("%Y-%m-%d")

    result_dict_1_9, monitoring_1_9_log = check_1_9()
    message_dict['full_log'] = {**monitoring_1_9_log}
    message_dict['details'] = {**result_dict_1_9}
    sorted_full_log = dict(sorted(message_dict['full_log'].items()))
    message_dict['full_log'] = sorted_full_log
    sorted_details = dict(sorted(message_dict['details'].items()))
    message_dict['details'] = sorted_details

    if result_dict_1_9:
        message_dict['status'] = "FAILED"
        message_dict['info'] = "Some steps failed"
    else:
        message_dict['status'] = "OK"
        message_dict['info'] = "All steps executed without errors"

    return message_dict


if __name__ == '__main__':
    # # check command line parameters and print or email result
    parser = create_arg_parser()
    arg_list = parser.parse_args()
    if len(sys.argv[1:]) > 0:
        result = run_steps()
        with open('/tmp/ae_retargeting_dashboard_daily.json', 'w') as json_file:
            json.dump(result, json_file, indent=4, default=str)

        if arg_list.recipients:
            recipients = arg_list.recipients[0]
            send_mail(recipients=recipients, message=json.dumps(result, indent=4, default=str), status=result['status'])

        if arg_list.print:
            print(json.dumps(result, indent=4, default=str))
    else:
        parser.print_help()