"""User delete request asset for processing RDS user deletion request data and triggering Glue export."""

from datetime import datetime
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_24H_FAIL_25H, 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 daily_cron_condition

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

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

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


@dg.asset(
    name="rds_user_delete_request",
    description="User delete request table with data from RDS (daily partitioned).",
    group_name=Group.RDS.value,
    partitions_def=dg.DailyPartitionsDefinition(
        start_date=USER_DELETE_REQUEST_START_DATE, 
        end_offset=0
    ),
    backfill_policy=dg.BackfillPolicy.multi_run(max_partitions_per_run=7),
    owners=[Team.CORE_POD.value],
    metadata={
        "database": SNOWFLAKE_DB,
        "schema": SNOWFLAKE_SCHEMA,
        "table_name": USER_DELETE_REQUEST_TABLE_NAME,
        "data_start_date": USER_DELETE_REQUEST_START_DATE.strftime("%Y-%m-%d"),
        "cluster_by": "[id]",
        "partition_expr": PartitionExpr.DAILY.value,
        "transient": True,
        "sla_minutes": 60,
    },
    automation_condition=daily_cron_condition,
    freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_24H_FAIL_25H,
)
def rds_user_delete_request(
    context: dg.AssetExecutionContext, 
    snowflake: SnowflakeResource
) -> dg.MaterializeResult:
    """
    Materializes the RDS_USER_DELETE_REQUEST table in Snowflake.
    
    This asset performs the following steps:
    1. Triggers an AWS Glue job to export the latest user delete request data from RDS to S3.
    2. Sets the appropriate Snowflake warehouse.
    3. Creates the target table if it does not already exist.
    4. Deletes existing data for the current daily partition in Snowflake.
    5. Upserts the new data from the S3 staging path into the Snowflake table.
    
    The asset is partitioned daily and includes data quality checks for ID uniqueness and row count.
    """
    run_id = context.run.run_id
    logger = dg.get_dagster_logger()
    partition_start = context.partition_time_window.start
    partition_end = context.partition_time_window.end
    
    logger.info(f"Processing partition: {partition_start} to {partition_end}")
    
    # Prepare SQL parameters
    fetch_params = {
        "partition_start_date": partition_start.strftime("%Y-%m-%d"),
        "partition_end_date": partition_end.strftime("%Y-%m-%d"),
        "user_delete_request_table_name": USER_DELETE_REQUEST_TABLE_NAME,
        "stage_path": f"@SUNO_DATABASE_EVENTS/bots_userdeleterequest/pdate={partition_start.strftime('%Y-%m-%d')}",
    }
    
    rows_affected = 0
    
    try:
        # Step 1: Trigger Glue job to export from RDS to S3
        logger.info("Step 1: Triggering Glue job to export data from RDS to S3")
        glue_result = trigger_glue_job(
            job_name="rds_to_s3_bots_userdeleterequest_daily_upsert",
            context=context,
            arguments={
                "--partition_date": partition_start.strftime("%Y-%m-%d"),
            },
            poll_interval=20,
            wait_for_completion=True,
            timeout=1800  # 30 minutes timeout
        )
        
        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
        raise e
    
    with snowflake.get_connection() as conn:
        cursor = conn.cursor()
        
        # Step 2: Set warehouse
        logger.info("Step 2: Use Warehouse")
        use_warehouse_query = load_query("src/utils/snowflake/queries/use_warehouse.sql", params={"warehouse": WAREHOUSE})
        log_query(logger, use_warehouse_query)
        cursor.execute(use_warehouse_query)
        
        
        # Step 4: Delete existing partition data
        logger.info("Step 4: Delete Daily Partition")
        delete_partition_query = load_query("src/utils/snowflake/queries/delete_daily_partitions.sql", params={**fetch_params, "delete_partition_table_name": USER_DELETE_REQUEST_TABLE_NAME})
        log_query(logger, delete_partition_query)
        cursor.execute(delete_partition_query)
        
        # Step 5: Upsert data from S3
        logger.info("Step 5: Upsert Data")
        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 upserted data into {USER_DELETE_REQUEST_TABLE_NAME}. Affected {rows_affected} rows.")
    
    return dg.MaterializeResult(
        metadata={
            "run_id": dg.MetadataValue.text(run_id),
            "table_name": USER_DELETE_REQUEST_TABLE_NAME,
            "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,
        }
    )