from datetime import datetime
import uuid

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

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

# Hardcoded date range for backfill
start_date = datetime(2024, 1, 1)  # Update this to your desired start date
end_date = datetime(2024, 1, 31)    # Update this to your desired end date

APPLICATION_NAME = "snowflake_to_rds_stem_clip_backfill"

insert_sql = """
INSERT INTO bots_steminfo 
(id, created_at, updated_at, clip_from_id, task, stem_task, 
    stem_type_id, stem_type_group_name, control_tags, clip_id) 
VALUES %s
ON CONFLICT (clip_id) 
DO UPDATE SET 
    updated_at = EXCLUDED.updated_at,
    clip_from_id = EXCLUDED.clip_from_id,
    task = EXCLUDED.task,
    stem_task = EXCLUDED.stem_task,
    stem_type_id = EXCLUDED.stem_type_id,
    stem_type_group_name = EXCLUDED.stem_type_group_name,
    control_tags = EXCLUDED.control_tags
"""
    
snowflake_query_sql = """
SELECT 
    *
FROM stem_table
WHERE p_date = '{process_date}'
"""

def stem_clip_row_mapper(row):
    """Convert Snowflake row to tuple for bots_steminfo"""
    # Clean JSON-encoded strings (remove extra quotes)
    def clean_value(val):
        if val is None:
            return None
        if isinstance(val, str) and val.startswith('"') and val.endswith('"'):
            return val[1:-1]  # Remove surrounding quotes
        return val
    
    return (
        str(uuid.uuid4()),                    # id
        row.CREATED_AT,                       # created_at
        row.UPDATED_AT,                       # updated_at  
        clean_value(row.CLIP_FROM_ID),        # clip_from_id
        clean_value(row.TASK),                # task (remove quotes)
        clean_value(row.STEM_TASK),           # stem_task (remove quotes)
        row.STEM_TYPE_ID if row.STEM_TYPE_ID != 'null' else None, # stem_type_id
        row.STEM_TYPE_GROUP_NAME,             # stem_type_group_name
        row.CONTROL_TAGS,                     # control_tags
        row.CLIP_ID,                          # clip_id (matches INSERT SQL order)
    )

# Filter batch to only include clip_ids that exist in target_table
def filter_invalid_rows(batch, read_pg_cursor):
    clip_ids_in_batch = [row[9] for row in batch if row[9] is not None]  # clip_id is at index 9
    
    if clip_ids_in_batch:
        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()}
        # Return filtered batch instead of just existing_clip_ids
        filtered_batch = [row for row in batch if row[9] in existing_clip_ids]
        return filtered_batch
    return []

backfill_data_day_by_day_from_snowflake(spark, start_date, end_date, snowflake_query_sql, insert_sql, stem_clip_row_mapper, APPLICATION_NAME, filter_invalid_rows)

job.commit()