"""Clip embedding asset for processing RDS clip embedding data and triggering Glue export."""

from datetime import datetime, timedelta
import warnings
from pathlib import Path

import dagster as dg
from dagster_snowflake import SnowflakeResource

from src.utils.snowflake.constants import TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H, Group, PartitionExpr, Warehouse, SnowflakeDB, Team, SnowflakeSchema
from src.utils.snowflake.logger import log_query
from src.utils.snowflake.query import load_query
from src.utils.glue_utils import trigger_glue_job
from src.utils.automation_conditions import rds_hourly_cron_condition

warnings.filterwarnings("ignore", category=dg.BetaWarning)

# Get directory of this file for relative SQL file loading
ASSET_DIR = Path(__file__).parent

CLIP_EMBEDDING_START_DATE = datetime.strptime('2024-01-01', '%Y-%m-%d')
CLIP_EMBEDDING_TABLE_NAME = "RDS_CLIP_EMBEDDING"
WAREHOUSE = Warehouse.SUNO_PROD_RDS_HOURLY_X_SMALL.value
SNOWFLAKE_DB = SnowflakeDB.SUNO_PROD.value
SNOWFLAKE_SCHEMA = SnowflakeSchema.PROD.value


@dg.asset(
    name="rds_clip_embedding",
    description="Processed clip embedding table with vectors and similarity data from RDS, triggers Glue export to S3.",
    group_name=Group.RDS.value,
    partitions_def=dg.HourlyPartitionsDefinition(start_date=CLIP_EMBEDDING_START_DATE, end_offset=0),
    backfill_policy=dg.BackfillPolicy.multi_run(max_partitions_per_run=24),
    owners=[Team.CORE_POD.value],
    metadata={
        "database": SNOWFLAKE_DB,
        "schema": SNOWFLAKE_SCHEMA,
        "table_name": CLIP_EMBEDDING_TABLE_NAME,
        "data_start_date": CLIP_EMBEDDING_START_DATE.strftime("%Y-%m-%d"),
        "cluster_by": "[p_date, p_hour]",
        "partition_expr": PartitionExpr.HOURLY.value,
        "transient": True,
        "sla_minutes": 60,
    },
    automation_condition=rds_hourly_cron_condition,
    freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H,
)
def rds_clip_embedding(
    context: dg.AssetExecutionContext,
    snowflake: SnowflakeResource
) -> dg.MaterializeResult:
    """
    Process clip embedding data from RDS and trigger Glue export.

    Steps:
    1. Trigger Glue job to export to S3
    3. Delete existing data for partition
    4. Upsert data from S3
    """
    run_id = context.run.run_id
    logger = dg.get_dagster_logger()

    # Get partition time window for processing
    partition_start = context.partition_time_window.start
    partition_end = context.partition_time_window.end

    fetch_params = {
        "partition_start_date": partition_start.strftime("%Y-%m-%d"),
        "partition_end_date": partition_end.strftime("%Y-%m-%d"),
        "partition_start_hour": partition_start.hour,
        "partition_end_hour": partition_end.hour,
        "clip_embedding_table_name": CLIP_EMBEDDING_TABLE_NAME,
        "stage_path": f"@SUNO_DATABASE_EVENTS/bots_clipembedding/pdate={partition_start.strftime('%Y-%m-%d')}/phour={partition_start.strftime('%H')}",
    }

    logger.info(f"Processing rds_clip_embedding for partition: {partition_start} to {partition_end}")
    logger.info(f"Fetch params: {fetch_params}")

    rows_affected = 0

    # Step 1: Trigger Glue job to export data to S3
    logger.info("Step 1: Triggering Glue job to export clip embedding data to S3...")
    try:
        glue_result = trigger_glue_job(
            job_name="rds_to_s3_bots_clipembedding_hourly_upsert",
            context=context,
            arguments={
                "--partition_date": partition_start.strftime("%Y-%m-%d"),
                "--partition_hour": str(partition_start.hour),
            },
            poll_interval=20,
            wait_for_completion=True,
            timeout=1800,  # 30 minutes
        )
        logger.info(f"Glue job completed successfully: {glue_result}")
        glue_job_status = "SUCCESS"
        glue_job_id = glue_result.get("job_run_id", "N/A")
        glue_execution_time = glue_result.get("execution_time", 0)
    except Exception as e:
        logger.error(f"Glue job failed: {str(e)}")
        glue_job_status = "FAILED"
        glue_job_id = "N/A"
        glue_execution_time = 0
        # Re-raise the exception to fail the asset materialization
        raise e

    # Step 2: Process data in Snowflake
    with snowflake.get_connection() as conn:
        cursor = conn.cursor()
        logger.info(f"Using warehouse {WAREHOUSE}")
        warehouse_query = load_query(
            "src/utils/snowflake/queries/use_warehouse.sql",
            params={"warehouse": WAREHOUSE}
        )
        log_query(logger, warehouse_query)
        cursor.execute(warehouse_query)


        # Delete existing data for the partition window
        logger.info(f"Deleting existing data from {CLIP_EMBEDDING_TABLE_NAME}...")
        delete_query = load_query(
            "src/utils/snowflake/queries/delete_hourly_partitions.sql",
            params={**fetch_params, "delete_partition_table_name": CLIP_EMBEDDING_TABLE_NAME}
        )
        log_query(logger, delete_query)
        cursor.execute(delete_query)

        # Upsert data from S3
        logger.info(f"Upserting data into {CLIP_EMBEDDING_TABLE_NAME} from S3.")
        upsert_query = load_query(
            ASSET_DIR / "upsert.sql",
            params=fetch_params
        )
        log_query(logger, upsert_query)
        cursor.execute(upsert_query)

        rows_affected = cursor.rowcount

    logger.info(f"Successfully processed partition. Affected {rows_affected} rows.")

    return dg.MaterializeResult(
        metadata={
            "run_id": dg.MetadataValue.text(run_id),
            "table_name": CLIP_EMBEDDING_TABLE_NAME,
            "partition_time_window_start": dg.MetadataValue.text(partition_start.isoformat()),
            "partition_time_window_end": dg.MetadataValue.text(partition_end.isoformat()),
            "dagster/row_count": rows_affected,
            "glue_job_status": dg.MetadataValue.text(glue_job_status),
            "glue_job_run_id": dg.MetadataValue.text(glue_job_id),
            "glue_execution_time_seconds": glue_execution_time,
        },
    )