"""ML Song Summary Info Hourly asset for aggregating song play metrics."""

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.automation_conditions import hourly_cron_with_eager_historical_backfill_condition

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

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

ML_SONG_SUMMARY_START_DATE = datetime.strptime('2024-01-01', '%Y-%m-%d')
ML_SONG_SUMMARY_TABLE_NAME = "ML_SONG_SUMMARY_INFO"
WAREHOUSE = Warehouse.SUNO_PROD_ML_HOURLY_LARGE.value
SNOWFLAKE_DB = SnowflakeDB.SUNO_PROD.value
SNOWFLAKE_SCHEMA = SnowflakeSchema.PROD.value


@dg.asset(
    name="ml_song_summary_info_hourly_insert",
    description="Aggregates song play metrics on an hourly basis by combining current hour play data with previous cumulative data. Depends on agg_play_info_hourly_v0_insert.",
    group_name=Group.AGG.value,
    partitions_def=dg.HourlyPartitionsDefinition(
        start_date=ML_SONG_SUMMARY_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": ML_SONG_SUMMARY_TABLE_NAME,
        "data_start_date": ML_SONG_SUMMARY_START_DATE.strftime("%Y-%m-%d"),
        "partition_expr": PartitionExpr.HOURLY.value,
        "sla_minutes": 60,
        "depends_on": "agg_play_info_hourly_v0_insert",
    },
    deps=["agg_play_info_hourly_v0_insert"],  # This asset depends on agg_play_info_hourly_v0_insert
    automation_condition=hourly_cron_with_eager_historical_backfill_condition,
    freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H,
)
def ml_song_summary_info_hourly_insert(
    context: dg.AssetExecutionContext, 
    snowflake: SnowflakeResource
) -> dg.MaterializeResult:
    """
    Aggregate song play metrics hourly by combining new data with previous cumulative data.
    
    Steps:
    1. Set warehouse
    2. Calculate current and previous partition date/hour
    3. Insert aggregated song summary data
    4. Call validation stored procedure
    """
    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
    
    # Calculate previous partition (1 hour before)
    previous_partition = partition_start - timedelta(hours=1)

    fetch_params = {
        "p_date": partition_start.strftime("%Y-%m-%d"),
        "p_hour": partition_start.hour,
        "pre_p_date": previous_partition.strftime("%Y-%m-%d"),
        "pre_p_hour": previous_partition.hour,
    }

    logger.info(f"Processing ml_song_summary_info for partition: {partition_start} to {partition_end}")
    logger.info(f"Current partition: {fetch_params['p_date']} hour {fetch_params['p_hour']}")
    logger.info(f"Previous partition: {fetch_params['pre_p_date']} hour {fetch_params['pre_p_hour']}")

    rows_affected = 0

    # Process data in Snowflake
    with snowflake.get_connection() as conn:
        cursor = conn.cursor()
        
        # Step 1: Set warehouse
        logger.info(f"Step 1: 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)
    
        # Step 2: Insert aggregated song summary data
        logger.info(f"Step 2: Inserting aggregated song summary data into {ML_SONG_SUMMARY_TABLE_NAME}...")
        insert_query = load_query(
            ASSET_DIR / "insert.sql", 
            params=fetch_params
        )
        log_query(logger, insert_query)
        cursor.execute(insert_query)
        rows_affected = cursor.rowcount
        logger.info(f"Inserted {rows_affected} rows")
        
        # Step 3: Call validation stored procedure
        logger.info(f"Step 3: Calling ML_SONG_SUMMARY_INFO_VALIDATION procedure...")
        validate_query = load_query(
            ASSET_DIR / "validate.sql", 
            params=fetch_params
        )
        log_query(logger, validate_query)
        cursor.execute(validate_query)
        
        logger.info("Validation procedure completed successfully")

    logger.info(f"Successfully processed {ML_SONG_SUMMARY_TABLE_NAME}. Total rows inserted: {rows_affected}")

    return dg.MaterializeResult(
        metadata={
            "run_id": dg.MetadataValue.text(run_id),
            "table_name": ML_SONG_SUMMARY_TABLE_NAME,
            "partition_time_window_start": dg.MetadataValue.text(partition_start.isoformat()),
            "partition_time_window_end": dg.MetadataValue.text(partition_end.isoformat()),
            "partition_date": dg.MetadataValue.text(partition_start.strftime("%Y-%m-%d")),
            "partition_hour": dg.MetadataValue.int(partition_start.hour),
            "previous_partition_date": dg.MetadataValue.text(previous_partition.strftime("%Y-%m-%d")),
            "previous_partition_hour": dg.MetadataValue.int(previous_partition.hour),
            "dagster/row_count": rows_affected,
        },
    )