from datetime import datetime

from utils.util import init_spark_context
from utils.postgres.backfill import update_data_day_by_day_from_snowflake

# Glue parameters
sc, glueContext, spark, job = init_spark_context()

# Hardcoded date range for backfill
start_date = datetime(2025, 8, 1)  # August 1, 2025
end_date = datetime(2025, 9, 24)  # September 24, 2025

APPLICATION_NAME = "remix_clip_backfill"

# Update SQL to set is_remix = true in metadata for clips with task = 'cover' or 'extend'
update_sql = """
UPDATE bots_generatedclip
SET metadata = jsonb_set(
    metadata,
    '{{is_remix}}',
    'true'::jsonb
)
WHERE id in ({format_strings})
"""

# Query to get clips that need to be updated
snowflake_query_sql = """
SELECT 
    id
FROM backfill_clip_remix
WHERE p_date = '{process_date}'
"""


def remix_clip_row_mapper(row):
    return row.ID,  # clip_id for WHERE clause    


def filter_invalid_rows(batch, read_pg_cursor):
    """Filter to only include clip_ids that exist in bots_generatedclip and return filtered batch"""
    clip_ids_in_batch = [row[0] for row in batch if row[0] is not None]
    
    if not clip_ids_in_batch:
        print("No valid clip IDs in batch")
        return []

    format_strings = ",".join(["%s"] * len(clip_ids_in_batch))
    read_pg_cursor.execute(
        f"""
        SELECT id FROM bots_generatedclip WHERE id IN ({format_strings})
    """,
        clip_ids_in_batch,
    )
    existing_clip_ids = {row[0] for row in read_pg_cursor.fetchall()}
    print(f"Found {len(existing_clip_ids)} existing clip IDs out of {len(clip_ids_in_batch)} total")
    
    # Return the filtered batch directly
    filtered_batch = [row for row in batch if row[0] in existing_clip_ids]
    print(f"Filtered batch from {len(batch)} to {len(filtered_batch)} rows")
    return filtered_batch


update_data_day_by_day_from_snowflake(
    spark, start_date, end_date, snowflake_query_sql, update_sql, remix_clip_row_mapper, APPLICATION_NAME, filter_invalid_rows
)

job.commit()