"""Aggregation asset for AGG_PLAY_INFO_HOURLY_V0 table."""

from datetime import datetime
import warnings
from pathlib import Path

import dagster as dg
from dagster_snowflake import SnowflakeResource

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

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

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

AGG_PLAY_INFO_START_DATE = datetime.strptime('2024-01-01', '%Y-%m-%d')
AGG_PLAY_INFO_TABLE_NAME = "AGG_PLAY_INFO_HOURLY_V0"
WAREHOUSE = Warehouse.SUNO_PROD_AGG_HOURLY_LARGE.value
SNOWFLAKE_DB = SnowflakeDB.SUNO_PROD.value
SNOWFLAKE_SCHEMA = SnowflakeSchema.PROD.value


@dg.asset(
    name="agg_play_info_hourly_v0_insert",
    description="Aggregates play information on an hourly basis by calling agg_play_info_hourly_v0_proc stored procedure. Depends on rds_discord_info.",
    group_name=Group.AGG.value,
    partitions_def=dg.HourlyPartitionsDefinition(start_date=AGG_PLAY_INFO_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": AGG_PLAY_INFO_TABLE_NAME,
        "procedure_name": "agg_play_info_hourly_v0_proc",
        "data_start_date": AGG_PLAY_INFO_START_DATE.strftime("%Y-%m-%d"),
        "partition_expr": PartitionExpr.HOURLY.value,
        "sla_minutes": 60,
    },
    deps=["rds_discord_info"],  # This asset depends on rds_discord_info
    automation_condition=hourly_cron_with_eager_historical_backfill_condition,
    freshness_policy=TIME_WINDOW_FRESHNESS_POLICY_WARN_1H_FAIL_2H,
)
def agg_play_info_hourly_v0_insert(
    context: dg.AssetExecutionContext, 
    snowflake: SnowflakeResource
) -> dg.MaterializeResult:
    """
    Call agg_play_info_hourly_v0_proc to aggregate play information hourly.
    
    Steps:
    1. Set warehouse
    2. Call agg_play_info_hourly_v0_proc with partition date and hour
    """
    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 = {
        "start_date": partition_start.strftime("%Y-%m-%d"),
        "start_hour": partition_start.hour,
    }

    logger.info(f"Processing agg_play_info_hourly_v0 for partition: {partition_start} to {partition_end}")
    logger.info(f"Calling procedure with params: {fetch_params}")

    rows_affected = 0

    # Process data in Snowflake
    with snowflake.get_connection() as conn:
        cursor = conn.cursor()
        
        # 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)
    
        # Call stored procedure
        logger.info(f"Step 2: Calling agg_play_info_hourly_v0_proc...")
        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"Successfully called agg_play_info_hourly_v0_proc. Rows affected: {rows_affected}")

    return dg.MaterializeResult(
        metadata={
            "run_id": dg.MetadataValue.text(run_id),
            "table_name": AGG_PLAY_INFO_TABLE_NAME,
            "procedure_name": "agg_play_info_hourly_v0_proc",
            "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),
            "dagster/row_count": rows_affected,
        },
    )