from datetime import timedelta
import time
import random

from utils.snowflake.utils import get_data_from_snowflake
from psycopg2.extras import execute_values
from utils.postgres.postgres_client import get_postgres_connection
from utils.util import get_environment
from typing import List, Callable, Optional

environment = get_environment()
database_name = f"SUNO_{environment}"


# Insert batch into target table
def insert_batch(batch, write_pg_cursor, write_pg_connection, read_pg_cursor, insert_sql, filter_invalid_rows_func):
    if filter_invalid_rows_func:
        original_count = len(batch)
        filtered_batch = filter_invalid_rows_func(batch, read_pg_cursor)
        batch[:] = filtered_batch
        print(f"Filtered batch from {original_count} to {len(batch)} rows (only clip_ids in bots_generatedclip)")
    else:
        print(f"No filter_invalid_rows_func provided, inserting all {len(batch)} rows")
    
    if not batch:
        print("No valid clip_ids found in batch, skipping insert")
        return
    
    for retry in range(3):
        sleep_time = 0.5 * (2**retry) + random.uniform(0, 0.1)
        time.sleep(sleep_time)
        try:
            execute_values(
                write_pg_cursor,
                insert_sql,
                batch,
            )
            write_pg_connection.commit()
            print(f"Inserted {len(batch)} rows")
            break
        except Exception as e:
            print(f"Error processing batch: {e}")
            continue
            

def backfill_data(spark, data, insert_sql, APPLICATION_NAME, row_mapping_func, filter_invalid_rows_func):    
    try:
        # Process data directly - no DataFrame needed, avoids schema inference issues
        print(f"Processing {len(data):,} rows directly from Snowflake")
        
        write_pg_connection = get_postgres_connection(APPLICATION_NAME)
        write_pg_cursor = write_pg_connection.cursor()
        read_pg_connection = get_postgres_connection(APPLICATION_NAME)
        read_pg_cursor = read_pg_connection.cursor()
        
        batch = []
        batch_size = 1000
        
        for row in data:
            try:
                # Use row mapping function to convert row to tuple
                mapped_row = row_mapping_func(row)
                batch.append(mapped_row)
                
                if len(batch) >= batch_size:
                    insert_batch(batch, write_pg_cursor, write_pg_connection, read_pg_cursor, insert_sql, filter_invalid_rows_func)
                    batch = []

            except Exception as e:
                print(f"Error processing row: {row}")
                raise e

        # Process remaining batch
        if batch:
            insert_batch(batch, write_pg_cursor, write_pg_connection, read_pg_cursor, insert_sql, filter_invalid_rows_func)
        
        # Clean up connections
        write_pg_cursor.close()
        write_pg_connection.close()
        read_pg_cursor.close()
        read_pg_connection.close()

        return len(data)
    except Exception as e:
        print(f"Error processing data: {str(e)}")
        raise e


# Backfill data day by day
def backfill_data_day_by_day_from_snowflake(spark, start_date, end_date, snowflake_query_sql, insert_sql, row_mapping_func, APPLICATION_NAME, filter_invalid_rows_func=None):
    total_processed = 0
    failed_dates = []
    try:
        print(f"Starting backfill from {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')}")
        current_date = start_date
        while current_date <= end_date:
            try:
                query_sql = snowflake_query_sql.format(process_date=current_date.strftime('%Y-%m-%d'))
                data = get_data_from_snowflake("SUNO_GLUE_BACKFILL_LARGE", database_name, environment, query_sql)
                processed_count = backfill_data(spark, data, insert_sql, APPLICATION_NAME, row_mapping_func, filter_invalid_rows_func)
                total_processed += processed_count  
                print(f"Successfully processed {processed_count} rows for {current_date}")
                if processed_count > 0:
                    time.sleep(1)
            except Exception as e:
                print(f"Failed to process {current_date.strftime('%Y-%m-%d')}: {str(e)}")
                failed_dates.append(current_date.strftime('%Y-%m-%d'))    
            current_date += timedelta(days=1)

        print(f"Total rows processed: {total_processed:,}")
        if failed_dates:
            print(f"Failed dates ({len(failed_dates)}): {', '.join(failed_dates)}")
            raise Exception(f"Failed dates ({len(failed_dates)}): {', '.join(failed_dates)}")
        else:
            print("All dates processed successfully!")
    except Exception as e:
        print(f"Critical error occurred: {str(e)}")
        raise e


# Update batch into target table
def update_batch(batch: List, write_pg_cursor, write_pg_connection, update_sql: str, filter_invalid_rows_func: Optional[Callable] = None):
    """
    Update a batch of records in PostgreSQL using raw SQL method
    
    Args:
        batch: List of data rows to update
        write_pg_cursor: PostgreSQL cursor for writing
        write_pg_connection: PostgreSQL connection for writing
        update_sql: SQL update statement with {format_strings} placeholder
        filter_invalid_rows_func: Optional function to filter rows before update
    """
    if filter_invalid_rows_func:
        read_pg_connection = get_postgres_connection("analystic-database-connection")
        read_pg_cursor = read_pg_connection.cursor()
        try:
            original_count = len(batch)
            batch = filter_invalid_rows_func(batch, read_pg_cursor)
            print(f"Filtered batch from {original_count} to {len(batch)} rows")
        finally:
            read_pg_cursor.close()
            read_pg_connection.close()
    else:
        print(f"No filter_invalid_rows_func provided, updating all {len(batch)} rows")
    
    if not batch:
        print("No valid records found in batch, skipping update")
        return
    
    for retry in range(3):
        sleep_time = 0.5 * (2**retry) + random.uniform(0, 0.1)
        time.sleep(sleep_time)
        try:
            # Create format_strings placeholder for IN clause using raw SQL method
            ids_in_batch = [row[0] for row in batch if row[0] is not None]
            quoted_ids = [f"'{id}'" for id in ids_in_batch]
            format_strings_raw = ",".join(quoted_ids)
            # Only replace the specific {format_strings} parameter, not all braces
            raw_sql = update_sql.format(format_strings=format_strings_raw)
            print(f"Executing raw SQL update with {len(ids_in_batch)} values")
            write_pg_cursor.execute(raw_sql)
            write_pg_connection.commit()
            updated_count = write_pg_cursor.rowcount if hasattr(write_pg_cursor, 'rowcount') else len(batch)
            print(f"Updated {updated_count} rows")
            break
        except Exception as e:
            write_pg_connection.rollback()
            if retry < 2:
                print(f"Update failed (attempt {retry + 1}/3): {str(e)}, retrying...")
            else:
                print(f"Update failed after 3 attempts: {str(e)}")
                raise e


# Update data from Snowflake to PostgreSQL
def update_data(spark, data: List, update_sql: str, APPLICATION_NAME: str, row_mapping_func: Callable, filter_invalid_rows_func: Optional[Callable] = None):
    """
    Update data from Snowflake to PostgreSQL in batches
    
    Args:
        spark: Spark session
        data: List of Snowflake data rows
        update_sql: SQL update statement
        APPLICATION_NAME: Application name for connection
        row_mapping_func: Function to convert Snowflake row to PostgreSQL format
        filter_invalid_rows_func: Optional function to filter invalid rows
        
    Returns:
        int: Number of rows processed
    """
    if not data or len(data) == 0:
        print("No data to update")
        return 0

    try:
        write_pg_connection = get_postgres_connection("analystic-database-connection")
        write_pg_cursor = write_pg_connection.cursor()

        # Set application name for monitoring
        write_pg_cursor.execute(f"SET application_name = '{APPLICATION_NAME}'")
        write_pg_connection.commit()

        batch = []
        batch_size = 1000

        for row in data:
            batch.append(row_mapping_func(row))
            if len(batch) >= batch_size:
                update_batch(batch, write_pg_cursor, write_pg_connection, update_sql, filter_invalid_rows_func)
                batch = []

        # Update remaining batch
        if batch:
            update_batch(batch, write_pg_cursor, write_pg_connection, update_sql, filter_invalid_rows_func)

        write_pg_cursor.close()
        write_pg_connection.close()

        return len(data)
    except Exception as e:
        print(f"Error processing update data: {str(e)}")
        raise e


# Update data day by day from Snowflake
def update_data_day_by_day_from_snowflake(
    spark, 
    start_date, 
    end_date, 
    snowflake_query_sql: str, 
    update_sql: str, 
    row_mapping_func: Callable, 
    APPLICATION_NAME: str, 
    filter_invalid_rows_func: Optional[Callable] = None
):
    """
    Update data day by day from Snowflake to PostgreSQL
    
    Args:
        spark: Spark session
        start_date: Start date for processing
        end_date: End date for processing
        snowflake_query_sql: Snowflake SQL query with {process_date} placeholder
        update_sql: PostgreSQL update SQL statement
        row_mapping_func: Function to convert Snowflake row to PostgreSQL format
        APPLICATION_NAME: Application name for connection monitoring
        filter_invalid_rows_func: Optional function to filter invalid rows
        
    Returns:
        int: Total number of rows processed
    """
    total_processed = 0
    failed_dates = []
    try:
        print(f"Starting day-by-day update from {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')}")
        current_date = start_date
        while current_date <= end_date:
            try:
                query_sql = snowflake_query_sql.format(process_date=current_date.strftime('%Y-%m-%d'))
                data = get_data_from_snowflake("SUNO_GLUE_BACKFILL_LARGE", database_name, environment, query_sql)
                processed_count = update_data(spark, data, update_sql, APPLICATION_NAME, row_mapping_func, filter_invalid_rows_func)
                total_processed += processed_count  
                print(f"Successfully processed {processed_count} rows for {current_date.strftime('%Y-%m-%d')}")
                if processed_count > 0:
                    time.sleep(1)
            except Exception as e:
                print(f"Failed to process {current_date.strftime('%Y-%m-%d')}: {str(e)}")
                failed_dates.append(current_date.strftime('%Y-%m-%d'))    
            current_date += timedelta(days=1)

        print(f"Total rows processed: {total_processed:,}")
        if failed_dates:
            print(f"Failed dates ({len(failed_dates)}): {', '.join(failed_dates)}")
            raise Exception(f"Failed dates ({len(failed_dates)}): {', '.join(failed_dates)}")
        else:
            print("All dates processed successfully!")
            
        return total_processed
    except Exception as e:
        print(f"Critical error occurred during day-by-day update: {str(e)}")
        raise e