import pandas as pd
from datetime import datetime
import pytz, asyncio
from utils import execute_query, get_cloudwatch_logger, log_exception
from scripts.mysql_queries import overdue_calc_query, closed_cancelled_settled_check
logger = get_cloudwatch_logger(stream_name="overdue_update")
local_tz = pytz.timezone("Asia/Kolkata")


def update_consumer_application_element(result_df, batch_size=3000):
    """
    Updates principal_due, total_principal_due, amount_over_due, total_amount_due, emi_paid & emi_pending in consumer_application_element,
    for each application_id in batches.
    """
    now = datetime.now()
    sql = """
        update consumer_application_element
        set principal_due = :principal_due, total_principal_due = :total_principal_due, amount_over_due = :amount_over_due,
            total_amount_due = :total_amount_due, emi_paid = :emi_paid, emi_pending = :emi_pending, updated_at = :updated_at
        where application_id = :application_id
    """
    
    # Prepare parameter dicts for each row in the DataFrame
    all_params = [{
            "application_id": row.application_id, "principal_due": row.principal_due, "total_principal_due": row.total_principal_due,
            "amount_over_due": row.amount_over_due, "total_amount_due": row.total_amount_due,
            "emi_paid": row.emi_paid, "emi_pending": row.emi_pending, "updated_at": now
        }
        for row in result_df.itertuples(index=False)]
    
    # Execute the update in batches
    for i in range(0, len(all_params), batch_size):
        batch = all_params[i:i + batch_size]
        execute_query(sql, batch)
    
    return True


async def function1(data: dict) -> dict:
    logger.send_log_event("INFO: Inside overdue function..")

    app_id_list = data.get("app_id_list", [])
    try:
        # Step 1: Get all applications
        if app_id_list:
            all_active_cases = app_id_list # if app_id_list provided in request then update only for those application_ids, else for all
        else:
            all_app_query = """
                select id from consumer_application where state in ('disbursed', 'nach registered') 
                and (deleted_at is null and is_deleted is null and merchant_id not in (73, 9421))
            """
            all_app_df = execute_query(all_app_query)
            all_active_cases = all_app_df['id'].tolist()

        if not all_active_cases:
            return {"status": "FAILED", "message": "No cases in disbursed/nach registered status"}
        
        logger.send_log_event(f"INFO: Fetched {len(all_active_cases)} cases.")

        # Step 2: Execute overdue query
        params = {'app_ids': tuple(all_active_cases)} # tuple(all_active_cases) tuple([8674, 10649])
        result_df = execute_query(overdue_calc_query, params=params)
        
        # if case is foreclosed or settled then update all values to 0
        foreclosed_settled_df = execute_query(closed_cancelled_settled_check, params={"application_ids": list(result_df['application_id'])})
        if isinstance(foreclosed_settled_df, pd.DataFrame) and not foreclosed_settled_df.empty:
            result_df = result_df.merge(foreclosed_settled_df, on='application_id', how='left')
        mask = (result_df['foreclosed'] == 1) | (result_df['settled'] == 1) | (result_df['app_status'].isin(['closed','cancelled']))
        cols = ['principal_due', 'total_principal_due', 'amount_over_due', 'total_amount_due', 'emi_pending']
        result_df.loc[mask, cols] = 0
        result_df.loc[mask, 'emi_paid'] = result_df.loc[mask, 'emi_pending']
        logger.send_log_event(f"INFO: Calculated Overdue.")

        # Step 3: Update consumer_application_element as per result_df
        now = datetime.now(local_tz).strftime('%Y-%m-%d %H:%M:%S')
        
        # Call the functions for updating consumer_application_element
        ce_update_result = update_consumer_application_element(result_df)
        if ce_update_result:
            logger.send_log_event(f"INFO: consumer_application_additional and consumer_application_element tables successfully updated for {len(result_df)} records.")
            return {
                "status": "SUCCESS",
                "message": f"Columns principal_due, total_principal_due, amount_over_due, total_amount_due, emi_paid, emi_pending updated in consumer_application_element table for {len(result_df)} records at {now}"
                }

        # If any of the updates failed, return error message
        return {"status": "ERROR", "message": "Something went wrong during the overdue update process."}

    except Exception as e:
        log_exception(logger, "Exception occurred during overdue calculation")
        return {"status": "ERROR", "message": str(e)}


# if __name__ == '__main__':
#     data = {"app_id_list" : [8674, 10649, 12903, 10190, 13677]}
#     print(asyncio.run(function1(data)))