"""Hook manual quality label asset for processing RDS hook manual quality label 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

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


@dg.asset(
    name="rds_hook_manual_quality_label",
    description="Processed hook manual quality label table with data from RDS, triggers Glue export to S3.",
    group_name=Group.RDS.value,
    partitions_def=dg.HourlyPartitionsDefinition(start_date=HOOK_MANUAL_QUALITY_LABEL_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": HOOK_MANUAL_QUALITY_LABEL_TABLE_NAME,
        "data_start_date": HOOK_MANUAL_QUALITY_LABEL_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_hook_manual_quality_label(
    context: dg.AssetExecutionContext, 
    snowflake: SnowflakeResource
) -> dg.MaterializeResult:
    """
    Process hook manual quality label 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,
        "hook_manual_quality_label_table_name": HOOK_MANUAL_QUALITY_LABEL_TABLE_NAME,
        "stage_path": f"@SUNO_DATABASE_EVENTS/video_hookmanualqualitylabel/pdate={partition_start.strftime('%Y-%m-%d')}/phour={partition_start.strftime('%H')}",
    }

    logger.info(f"Processing rds_hook_manual_quality_label 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 hook manual quality label data to S3...")
    try:
        glue_result = trigger_glue_job(
            job_name="rds_to_s3_hook_manual_quality_label_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 {HOOK_MANUAL_QUALITY_LABEL_TABLE_NAME}...")
        delete_query = load_query(
            "src/utils/snowflake/queries/delete_hourly_partitions.sql", 
            params={**fetch_params, "delete_partition_table_name": HOOK_MANUAL_QUALITY_LABEL_TABLE_NAME}
        )
        log_query(logger, delete_query)
        cursor.execute(delete_query)
    
        # Upsert data from S3
        logger.info(f"Upserting data into {HOOK_MANUAL_QUALITY_LABEL_TABLE_NAME} from S3.")
        upsert_query = load_query(
            ASSET_DIR / "upsert.sql", 
            params=fetch_params
        )
        log_query(logger, upsert_query)
        cursor.execute(upsert_query)
    
        # Commit the transaction
        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": HOOK_MANUAL_QUALITY_LABEL_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,
        },
    )