"""Profile follow asset for processing RDS profile follow 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.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

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


@dg.asset(
    name="rds_profile_follow",
    description="Profile follow table (daily partitioned)",
    group_name=Group.RDS.value,
    partitions_def=dg.DailyPartitionsDefinition(start_date=PROFILE_FOLLOW_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": PROFILE_FOLLOW_TABLE_NAME,
        "cluster_by": "[p_date]",
        "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_profile_follow(
    context: dg.AssetExecutionContext, 
    snowflake: SnowflakeResource
) -> dg.MaterializeResult:
    """
    Process RDS profile follow data for a specific daily partition.
    
    This asset:
    1. Triggers AWS Glue job to export data from RDS to S3
    2. Creates the Snowflake table if it doesn't exist
    3. Deletes existing data for the partition
    4. Loads new data from S3 into Snowflake
    
    Args:
        context: Dagster execution context with partition information
        snowflake: Snowflake resource for database operations
        
    Returns:
        MaterializeResult with metadata about the operation
    """
    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"),
        "profile_follow_table_name": PROFILE_FOLLOW_TABLE_NAME,
        "stage_path": f"@SUNO_DATABASE_EVENTS/bots_profilefollow/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_profilefollow_daily_replace",
            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")
        cursor.execute(
            load_query(
                "src/utils/snowflake/queries/use_warehouse.sql",
                params={"warehouse": WAREHOUSE}
            )
        )
        
        # Step 4: Delete existing partition data
        logger.info("Step 4: Delete Daily Partition")
        cursor.execute(
            load_query(
                "src/utils/snowflake/queries/delete_daily_partitions.sql",
                params={**fetch_params, "delete_partition_table_name": PROFILE_FOLLOW_TABLE_NAME}
            )
        )
        
        # Step 5: Upsert data from S3
        logger.info("Step 5: Upsert Data")
        cursor.execute(
            load_query(
                ASSET_DIR / "upsert.sql",
                params=fetch_params
            )
        )
        
        rows_affected = cursor.rowcount
        logger.info(f"Successfully upserted data into {PROFILE_FOLLOW_TABLE_NAME}. Affected {rows_affected} rows.")
    
    return dg.MaterializeResult(
        metadata={
            "run_id": dg.MetadataValue.text(run_id),
            "table_name": PROFILE_FOLLOW_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,
        }
    )