"""Student plan eligible user asset for processing RDS student eligibility 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_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

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


@dg.asset(
    name="rds_student_plan_eligible_user",
    description="Processed student plan eligible user table with data from RDS, triggers Glue export to S3.",
    group_name=Group.RDS.value,
    partitions_def=dg.HourlyPartitionsDefinition(start_date=STUDENT_PLAN_ELIGIBLE_USER_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": STUDENT_PLAN_ELIGIBLE_USER_TABLE_NAME,
        "data_start_date": STUDENT_PLAN_ELIGIBLE_USER_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_student_plan_eligible_user(
    context: dg.AssetExecutionContext, 
    snowflake: SnowflakeResource
) -> dg.MaterializeResult:
    """
    Process student plan eligible user data from RDS and trigger Glue export.
    
    Steps:
    1. Trigger Glue job to export data to S3
    3. Delete existing data for partition
    4. Upsert data from S3 stage
    """
    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,
        "student_plan_eligible_user_table_name": STUDENT_PLAN_ELIGIBLE_USER_TABLE_NAME,
        "stage_path": f"@SUNO_DATABASE_EVENTS/bots_studentplaneligibleuser/pdate={partition_start.strftime('%Y-%m-%d')}/phour={partition_start.strftime('%H')}",
    }

    logger.info(f"Processing student_plan_eligible_user 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 student plan eligible user data to S3...")
    try:
        glue_result = trigger_glue_job(
            job_name="rds_to_s3_bots_studentplaneligibleuser_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 {STUDENT_PLAN_ELIGIBLE_USER_TABLE_NAME}...")
        delete_query = load_query(
            "src/utils/snowflake/queries/delete_hourly_partitions.sql", 
            params={**fetch_params, "delete_partition_table_name": STUDENT_PLAN_ELIGIBLE_USER_TABLE_NAME}
        )
        log_query(logger, delete_query)
        cursor.execute(delete_query)
    
        # Upsert data from S3 stage
        logger.info(f"Upserting data into {STUDENT_PLAN_ELIGIBLE_USER_TABLE_NAME} from S3 stage.")
        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": STUDENT_PLAN_ELIGIBLE_USER_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,
        },
    )