#! /usr/bin/env python3
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from contextlib import closing
from logging import info
from sql import sql_queries
import snowflake.connector
import psycopg2
import smtplib
import re
import argparse
import sys
import json
import datetime
import boto3
from boto3 import client

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']
sf_database = config['snowflake']['database']

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"]

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] RTI Push Notifications (weekly)"
    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,
        database=sf_database
    )
    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_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 check_2_to_8() -> dict:
    result_dict = {}
    full_log = {}

# Check_2 Data consistency - event generation check
    tmp_sf_check_2 = snowflake_fetch_data(sql_queries.sf_event_generation_check_2)
    sf_check_2_result = tmp_sf_check_2[0]['count']
    print(f"CHECK_2: - {sf_check_2_result}")
    full_log["CHECK_2"] = {'CHECK_2_RESULT': sf_check_2_result}

    if sf_check_2_result == 0:
        result_dict["CHECK_2"] = {'CHECK_2_RESULT': 'Failed: Query returned ' + str(sf_check_2_result)}
        full_log["CHECK_2"] = result_dict["CHECK_2"]
    else:
        full_log["CHECK_2"] = {'CHECK_2_RESULT': 'OK: Query returned ' + str(sf_check_2_result) + ' count'}

# Check_3 Data consistency - data input check
    tmp_sf_check_3 = snowflake_fetch_data(sql_queries.sf_data_input_check_3)
    sf_check_3_result = tmp_sf_check_3[0]['count']
    print(f"CHECK_3: - {sf_check_3_result}")
    full_log["CHECK_3"] = {'CHECK_3_RESULT': sf_check_3_result}

    if sf_check_3_result == 0:
        result_dict["CHECK_3"] = {'CHECK_3_RESULT': 'Failed: Query returned ' + str(sf_check_3_result)}
        full_log["CHECK_3"] = result_dict["CHECK_3"]
    else:
        full_log["CHECK_3"] = {'CHECK_3_RESULT': 'OK: Query returned ' + str(sf_check_3_result) + ' count'}

# Check_7 Missing artist mapping monitoring, and insertion process
    tmp_sf_check_7 = snowflake_fetch_data(sql_queries.sf_mapping_insertion_check_7)
    if tmp_sf_check_7:
        full_log["CHECK_7"] = {}
        temp_result_dict = {}
        temp_result_dict["CHECK_7"] = {}
        for item in tmp_sf_check_7:
            isrc = item['TRACK_ISRC']
            artist_id = item['GRAS_ARTIST_ID']
            app_name = item['APP_NAME']
            full_log["CHECK_7"][artist_id] = {'APP_NAME': app_name, 'ISRC': isrc}
            temp_result_dict["CHECK_7"][artist_id] = {'APP_NAME': app_name, 'ISRC': isrc}
            print(f"'GRAS_ARTIST_ID': {artist_id}, 'APP_NAME': {app_name}, 'ISRC': {isrc}")
        if temp_result_dict["CHECK_7"]:
            result_dict["CHECK_7"] = temp_result_dict["CHECK_7"]
    else:
        full_log["CHECK_7"] = {'CHECK_7_RESULT': 'OK: Query returned 0 results'}

# Check_8 Events export

    s3 = boto3.resource('s3')
    bucket = s3.Bucket('prod-delphi-notifications')
    result = bucket.meta.client.list_objects(Bucket=bucket.name,Delimiter='/')
    list_dirs = []
    for d in result.get('CommonPrefixes'):
        dir = d.get('Prefix')
        list_dirs.append(dir)
    print(f"{list_dirs}")

    tmp_sf_check_8 = snowflake_fetch_data(sql_queries.sf_events_export_check_8)
    full_log["CHECK_8"] = {}
    temp_result_dict = {}
    temp_result_dict["CHECK_8"] = {}
    for item in tmp_sf_check_8:
        date = item['CREATED_AT'].strftime("%Y%m%d")
        result = 'date=' + str(date) + '/'
        print(f"{result}")
        if result in list_dirs:
            full_log["CHECK_8"][date] = {result: 'OK: folder found in S3 bucket'}
        else:
            temp_result_dict["CHECK_8"][date] = {result: 'Failed: folder NOT FOUND in S3 bucket'}
        if temp_result_dict["CHECK_8"]:
            result_dict["CHECK_8"] = temp_result_dict["CHECK_8"]
            full_log["CHECK_8"] = result_dict["CHECK_8"]


    return result_dict, full_log

################# run the all steps ######################
def run_steps() -> dict:
    message_dict = {'name':"[PROD] RTI Push Notifications (weekly)",'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, full_log = check_2_to_8()

    message_dict['full_log'] = {**full_log}
    message_dict['details'] = {**result_dict}

    if result_dict:
        postgres_update_failure_count(sql_queries.failure_count_increase_weekly)
        message_dict['status'] = "FAILED"
        message_dict['info'] = "Some steps failed"

    else:
        postgres_update_failure_count(sql_queries.failure_count_to_zero_weekly)
        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/rti_push_notifications_weekly_log.json', 'w') as json_file:
            json.dump(result, json_file, indent=4)

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

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