"""Usage plan asset for processing RDS usage plan data and triggering Glue export."""

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

USAGE_PLAN_TABLE_NAME = "RDS_USAGE_PLAN"
WAREHOUSE = Warehouse.SUNO_PROD_RDS_HOURLY_X_SMALL.value
SNOWFLAKE_DB = SnowflakeDB.SUNO_PROD.value
SNOWFLAKE_SCHEMA = SnowflakeSchema.PROD.value


@dg.asset(
    name="rds_usage_plan",
    description="Usage plan table - full snapshot replaced daily.",
    group_name=Group.RDS.value,
    freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_24H_FAIL_25H,
    owners=[Team.CORE_POD.value],
    metadata={
        "database": SNOWFLAKE_DB,
        "schema": SNOWFLAKE_SCHEMA,
        "table_name": USAGE_PLAN_TABLE_NAME,
        "transient": True,
        "sla_minutes": 60,
    }
)
def rds_usage_plan(
    context: dg.AssetExecutionContext, 
    snowflake: SnowflakeResource
) -> dg.MaterializeResult:
    """
    Materializes the RDS_USAGE_PLAN table in Snowflake.
    
    This asset performs the following steps:
    1. Triggers an AWS Glue job to export the latest usage plan data from RDS to S3.
    2. Sets the appropriate Snowflake warehouse.
    3. Replaces the entire table with the new snapshot using CREATE OR REPLACE.
    
    The asset contains a full snapshot (non-incremental, non-partitioned).
    """
    run_id = context.run.run_id
    logger = dg.get_dagster_logger()
    
    # Use current date instead of partition date
    current_date = datetime.now(timezone.utc).date()
    
    logger.info(f"Processing usage plan snapshot for {current_date}")
    
    # Prepare SQL parameters
    fetch_params = {
        "usage_plan_table_name": USAGE_PLAN_TABLE_NAME,
        "stage_path": f"@SUNO_DATABASE_EVENTS/bots_usageplan/pdate={current_date.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_usageplan_manually",
            context=context,
            arguments={
                "--partition_date": current_date.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 3: Create or replace table (full snapshot)
        logger.info("Step 3: Create or Replace Table")
        create_table_query = load_query(ASSET_DIR / "create.sql", params=fetch_params)
        log_query(logger, create_table_query)
        cursor.execute(create_table_query)
        rows_affected = cursor.rowcount
        logger.info(f"Successfully replaced table {USAGE_PLAN_TABLE_NAME}. Loaded {rows_affected} rows.")
    
    return dg.MaterializeResult(
        metadata={
            "run_id": dg.MetadataValue.text(run_id),
            "table_name": USAGE_PLAN_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,
        }
    )