#! /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']

db_name_slz = config['postgres_slz']["database"]
db_user_slz = config['postgres_slz']["user"]
db_password_slz = config['postgres_slz']["password"]
db_host_slz = config['postgres_slz']["host"]
db_port_slz = config['postgres_slz']["port"]

db_name_main = config['postgres_main']["database"]
db_user_main = config['postgres_main']["user"]
db_password_main = config['postgres_main']["password"]
db_host_main = config['postgres_main']["host"]
db_port_main = config['postgres_main']["port"]

local_db_name = config['failure_counter']["database"]
local_db_user = config['failure_counter']["user"]
local_db_password = config['failure_counter']["password"]
local_db_host = config['failure_counter']["host"]
local_db_port = config['failure_counter']["port"]

failure_count_increase = "UPDATE failed_count SET failed_count = failed_count + 1, last = 'FAILED', last_updated = NOW() WHERE count_id = 'bw_daily';"
failure_count_to_zero = "UPDATE failed_count SET failed_count = 0, last = 'OK', last_updated = NOW() WHERE count_id = 'bw_daily';"

def postgres_update_failure_count(sql_query):
    with closing(psycopg2.connect(database=local_db_name, user=local_db_user, password=local_db_password, host=local_db_host, port=local_db_port)) as connection:
        with connection.cursor() as cursor:
            cursor.execute(sql_query)
            connection.commit()

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] Spotify Juno (Exploration Area)"
    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 postgres_fetch_data_main(sql_query) -> list:
    result_list = []
    with closing(psycopg2.connect(database=db_name_main, user=db_user_main, password=db_password_main, host=db_host_main, port=db_port_main)) as connection:
        with connection.cursor() as cursor:
            cursor.execute(sql_query)
            columns = cursor.description
            rows = cursor.fetchall()
    for row in rows:
        tmp = {}
        for i in range(len(columns)):
            tmp[columns[i][0]] = row[i]
        result_list.append(tmp)
    return result_list

def postgres_fetch_data_slz(sql_query) -> list:
    result_list = []
    with closing(psycopg2.connect(database=db_name_slz, user=db_user_slz, password=db_password_slz, host=db_host_slz, port=db_port_slz)) as connection:
        with connection.cursor() as cursor:
            cursor.execute(sql_query)
            columns = cursor.description
            rows = cursor.fetchall()
    for row in rows:
        tmp = {}
        for i in range(len(columns)):
            tmp[columns[i][0]] = row[i]
        result_list.append(tmp)
    return result_list

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

    ### Check 1 ################################################
    tmp_sf_check_1 = snowflake_fetch_data(sql_queries.sf_check_1)
    cc_threshold = tmp_sf_check_1[0]['CONTEXT_COUNT_LAST_MONTH_THRESHOLD']
    print(f"CC_THRESHOLD: - {tmp_sf_check_1[0]['CONTEXT_COUNT_LAST_MONTH_THRESHOLD']}")
    #Check
    if tmp_sf_check_1:
        result_file = open("results/" + datetime.date.today().strftime("%F") + ".txt", "w")
        result_file.write(str(cc_threshold))
        result_file.close()
        if int(cc_threshold) < 400:
            full_log["CHECK_1"] = {'FAILED: CONTEXT_COUNT_LAST_MONTH_THRESHOLD': cc_threshold}
            result_dict["CHECK_1"] = full_log["CHECK_1"]
        else:
            full_log["CHECK_1"] = {'OK: CONTEXT_COUNT_LAST_MONTH_THRESHOLD': cc_threshold}
    else:
        result_dict["CHECK_1"] = {'FAILED': 'Query returned nothing'}

    ### Check 2 ################################################
    with open('results/' + datetime.date.today().strftime("%F") + '.txt') as f:
        threshold_from_prev_step = f.read()
    sf_check_2 = f"""select REPORT_DATE \
    from DELPHI_EXPLORATION.SYS.INBOUND_FILE_STATUS \
    where DSP = 'SPOTIFY' \
      and REPORT_NAME = 'STREAMS' \
      and REPORT_DATE between current_date() - 4 and current_date() - 2 \
      and STATUS in ('LOADED', 'L_OUTDATED') \
      and PROCESSING_COMPLETED = 1 \
    group by REPORT_DATE \
    having count(distinct LICENSOR, CONTEXT) > {threshold_from_prev_step} \
    order by REPORT_DATE desc;"""

    tmp_sf_check_2 = snowflake_fetch_data(sf_check_2)
    print(f"CHECK_2: - Query returned {len(tmp_sf_check_2)} rows")
    today_minus_2 = date.today() - timedelta(2)
    today_minus_3 = date.today() - timedelta(3)
    today_minus_4 = date.today() - timedelta(4)
    if len(tmp_sf_check_2) == 3:
        check_2_result_0 = tmp_sf_check_2[0]['REPORT_DATE']
        check_2_result_1 = tmp_sf_check_2[1]['REPORT_DATE']
        check_2_result_2 = tmp_sf_check_2[2]['REPORT_DATE']
        print(f"Dates: {check_2_result_0}, {check_2_result_1}, {check_2_result_2}")
    if len(tmp_sf_check_2) != 3:
        result_dict["CHECK_2"] = {'FAILED': 'Query did not return 3 rows'}
        full_log["CHECK_2"] = result_dict["CHECK_2"]
    elif today_minus_2 != check_2_result_0 or today_minus_3 != check_2_result_1 or today_minus_4 != check_2_result_2:
        result_dict["CHECK_2"] = {'FAILED': 'Returned dates are not meet the expected results'}
        full_log["CHECK_2"] = []
        for item in tmp_sf_check_2:
            full_log["CHECK_2"].append(item['REPORT_DATE'].strftime("%Y-%m-%d"))
            print(f"{item['REPORT_DATE']}")
    else:
        full_log["CHECK_2"] = []
        for item in tmp_sf_check_2:
            full_log["CHECK_2"].append(item['REPORT_DATE'].strftime("%Y-%m-%d"))

    ### Check 3 ################################################
    tmp_sf_check_3 = snowflake_fetch_data(sql_queries.sf_check_3)
    print(f"CHECK_3: - Query returned {len(tmp_sf_check_3)} rows")
    if len(tmp_sf_check_3) == 3:
        check_3_result_0 = tmp_sf_check_3[0]['DATE_ENABLED']
        check_3_result_1 = tmp_sf_check_3[1]['DATE_ENABLED']
        check_3_result_2 = tmp_sf_check_3[2]['DATE_ENABLED']
        print(f"Dates: {check_3_result_0}, {check_3_result_1}, {check_3_result_2}")
    if len(tmp_sf_check_3) != 3:
        result_dict["CHECK_3"] = {'FAILED': 'Query did not return 3 rows'}
        full_log["CHECK_3"] = result_dict["CHECK_3"]
    elif today_minus_2 != check_3_result_0 or today_minus_3 != check_3_result_1 or today_minus_4 != check_3_result_2:
        result_dict["CHECK_3"] = {'FAILED': 'Returned dates are not meet the expected results'}
        full_log["CHECK_3"] = []
        for item in tmp_sf_check_3:
            full_log["CHECK_3"].append(item['DATE_ENABLED'].strftime("%Y-%m-%d"))
            print(f"{item['DATE_ENABLED']}")
    else:
        full_log["CHECK_3"] = []
        for item in tmp_sf_check_3:
            full_log["CHECK_3"].append(item['DATE_ENABLED'].strftime("%Y-%m-%d"))


    ### Check 4 #############################################
    tmp_sf_check_4 = snowflake_fetch_data(sql_queries.sf_check_4)
    print(f"CHECK_4: {tmp_sf_check_4}")
    if tmp_sf_check_4:
        full_log["CHECK_4"] = {'FAILED': 'Query returned some results'}
        result_dict["CHECK_4"] = full_log["CHECK_4"]
    else:
        full_log["CHECK_4"] = {'OK': 'Query returned 0 results'}

    ### Check 5 ################################################
    tmp_sf_check_5 = snowflake_fetch_data(sql_queries.sf_check_5)
    print(f"CHECK_5: - Query returned {len(tmp_sf_check_5)} rows")
    if len(tmp_sf_check_5) == 3:
        check_5_result_0 = tmp_sf_check_5[0]['REPORT_DATE']
        check_5_result_1 = tmp_sf_check_5[1]['REPORT_DATE']
        check_5_result_2 = tmp_sf_check_5[2]['REPORT_DATE']
        print(f"Dates: {check_5_result_0}, {check_5_result_1}, {check_5_result_2}")
    if len(tmp_sf_check_5) != 3:
        result_dict["CHECK_5"] = {'FAILED': 'Query did not return 3 rows'}
        full_log["CHECK_5"] = result_dict["CHECK_5"]
    elif today_minus_2 != check_5_result_0 or today_minus_3 != check_5_result_1 or today_minus_4 != check_5_result_2:
        result_dict["CHECK_5"] = {'FAILED': 'Returned dates are not meet the expected results'}
        full_log["CHECK_5"] = []
        for item in tmp_sf_check_5:
            full_log["CHECK_5"].append(item['REPORT_DATE'].strftime("%Y-%m-%d"))
            print(f"{item['REPORT_DATE']}")
    else:
        full_log["CHECK_5"] = []
        for item in tmp_sf_check_5:
            full_log["CHECK_5"].append(item['REPORT_DATE'].strftime("%Y-%m-%d"))

    ### Check 6 ################################################
    tmp_sf_check_6 = snowflake_fetch_data(sql_queries.sf_check_6)
    print(f"CHECK_6: - Query returned {len(tmp_sf_check_6)} rows")
    if len(tmp_sf_check_6) == 3:
        check_6_result_0 = tmp_sf_check_6[0]['REPORT_DATE']
        check_6_result_1 = tmp_sf_check_6[1]['REPORT_DATE']
        check_6_result_2 = tmp_sf_check_6[2]['REPORT_DATE']
        print(f"Dates: {check_6_result_0}, {check_6_result_1}, {check_6_result_2}")
    if len(tmp_sf_check_6) != 3:
        result_dict["CHECK_6"] = {'FAILED': 'Query did not return 3 rows'}
        full_log["CHECK_6"] = result_dict["CHECK_6"]
    elif today_minus_2 != check_6_result_0 or today_minus_3 != check_6_result_1 or today_minus_4 != check_6_result_2:
        result_dict["CHECK_6"] = {'FAILED': 'Returned dates are not meet the expected results'}
        full_log["CHECK_6"] = []
        for item in tmp_sf_check_6:
            full_log["CHECK_6"].append(item['REPORT_DATE'].strftime("%Y-%m-%d"))
            print(f"{item['REPORT_DATE']}")
    else:
        full_log["CHECK_6"] = []
        for item in tmp_sf_check_6:
            full_log["CHECK_6"].append(item['REPORT_DATE'].strftime("%Y-%m-%d"))

    ### Check 7 ################################################
    tmp_sf_check_7 = snowflake_fetch_data(sql_queries.sf_check_7)
    print(f"CHECK_7: {tmp_sf_check_7}")
    if tmp_sf_check_7 and len(tmp_sf_check_7) == 1:
        for item in tmp_sf_check_7:
            if item['LATEST_SOURCE_KEY_AVAILABLE'] != item['RAW_LATEST_KEY'] or item['LATEST_SOURCE_KEY_AVAILABLE'] != item['MAIN_LATEST_KEY'] or item['RAW_LATEST_KEY'] != item['MAIN_LATEST_KEY']:
                lska = item['LATEST_SOURCE_KEY_AVAILABLE']
                rlk = item['RAW_LATEST_KEY']
                mlk = item['MAIN_LATEST_KEY']
                full_log["CHECK_7"] = {'LATEST_SOURCE_KEY_AVAILABLE': lska, 'RAW_LATEST_KEY': rlk, 'MAIN_LATEST_KEY': mlk}
                result_dict["CHECK_7"] = full_log["CHECK_7"]
            else:
                for item in tmp_sf_check_7:
                    lska = item['LATEST_SOURCE_KEY_AVAILABLE']
                    rlk = item['RAW_LATEST_KEY']
                    mlk = item['MAIN_LATEST_KEY']
                    full_log["CHECK_7"] = {'LATEST_SOURCE_KEY_AVAILABLE': lska, 'RAW_LATEST_KEY': rlk, 'MAIN_LATEST_KEY': mlk}
    else:
        result_dict["CHECK_7"] = {'FAILED': 'Query returned more then 1 row'}
        full_log["CHECK_7"] = result_dict["CHECK_7"]

    ### Check 8 #############################################
    tmp_sf_check_8 = snowflake_fetch_data(sql_queries.sf_check_8)
    print(f"CHECK_8: {tmp_sf_check_8}")
    sf_check_8_result=int(tmp_sf_check_8[0]['DAYS_AFTER_LAST_REFRESH'])
    if sf_check_8_result >= 200:
        result_dict["CHECK_8"] = {'FAILED: DAYS_AFTER_LAST_REFRESH': sf_check_8_result}
        full_log["CHECK_8"] = result_dict["CHECK_8"]
    else:
        full_log["CHECK_8"] = {'OK: DAYS_AFTER_LAST_REFRESH': sf_check_8_result}

    return result_dict, full_log

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

    result_dict_1_to_8, monitoring_1_to_8_log = check_1_to_8()
    message_dict['full_log'] = {**monitoring_1_to_8_log}
    message_dict['details'] = {**result_dict_1_to_8}

    if result_dict_1_to_8:
        postgres_update_failure_count(failure_count_increase)
        message_dict['status'] = "FAILED"
        message_dict['info'] = "Some steps failed"
    else:
        postgres_update_failure_count(failure_count_to_zero)
        message_dict['status'] = "OK"
        message_dict['info'] = "All steps completed successfully"

    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/spotify_juno_daily_log.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()