"""Asset checks for bk_realtime_event to validate data quality."""

import dagster as dg
from dagster_snowflake import SnowflakeResource
from src.assets.snowflake.raw_table.backend.realtime.assets import BK_REALTIME_EVENT_TABLE_NAME, WAREHOUSE
from src.utils.snowflake.partition_utils import parse_hourly_partition_key


# Constants for validation
MIN_ROW_COUNT = 1  # Ensure at least 1 row per partition
MAX_VOLUME_CHANGE_RATE = 0.10  # Maximum 10% change from previous day


@dg.asset_check(asset="bk_realtime_event", blocking=False)
def check_minimum_row_count(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Ensure minimum row count threshold is met (at least 1 row)."""
    partition_key = context.op_execution_context.partition_key
    p_date, p_hour = parse_hourly_partition_key(partition_key)
    
    with snowflake.get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(f"USE WAREHOUSE {WAREHOUSE}")
        query = f"""
            SELECT COUNT(*) as row_count
            FROM {BK_REALTIME_EVENT_TABLE_NAME}
            WHERE p_date = '{p_date}' AND p_hour = {p_hour}
        """
        result = cursor.execute(query).fetchone()
        row_count = result[0] if result else 0
    
    passed = row_count >= MIN_ROW_COUNT
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "partition_key": partition_key,
            "p_date": p_date,
            "p_hour": p_hour,
            "row_count": row_count,
            "min_threshold": MIN_ROW_COUNT,
        },
        description=f"Row count: {row_count:,} (threshold: {MIN_ROW_COUNT:,})"
    )


@dg.asset_check(asset="bk_realtime_event", blocking=False)
def check_day_over_day_volume(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check that daily data volume doesn't decrease more than 20% from previous day."""
    partition_key = context.op_execution_context.partition_key
    p_date, p_hour = parse_hourly_partition_key(partition_key)
    
    # Maximum allowed decrease rate (20%)
    MAX_DECREASE_RATE = 0.20
    
    with snowflake.get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(f"USE WAREHOUSE {WAREHOUSE}")
        
        # Get both today's and yesterday's count in one query
        query = f"""
            SELECT 
                SUM(CASE WHEN p_date = '{p_date}' THEN 1 ELSE 0 END) as today_count,
                SUM(CASE WHEN p_date = DATEADD(day, -1, '{p_date}'::date) THEN 1 ELSE 0 END) as yesterday_count
            FROM {BK_REALTIME_EVENT_TABLE_NAME}
            WHERE p_date IN ('{p_date}', DATEADD(day, -1, '{p_date}'::date))
              AND p_hour = {p_hour}
        """
        result = cursor.execute(query).fetchone()
        today_count = result[0] if result and result[0] else 0
        yesterday_count = result[1] if result and result[1] else 0
    
    # Calculate change
    if yesterday_count == 0:
        # If no data yesterday, skip this check
        return dg.AssetCheckResult(
            passed=True,
            metadata={
                "partition_key": partition_key,
                "p_date": p_date,
                "p_hour": p_hour,
                "today_count": today_count,
                "yesterday_count": yesterday_count,
                "change_rate": None,
            },
            description=f"No data from yesterday to compare. Today: {today_count:,}"
        )
    
    # Check if volume increased or stayed same - always pass
    if today_count >= yesterday_count:
        change_rate = (today_count - yesterday_count) / yesterday_count
        return dg.AssetCheckResult(
            passed=True,
            metadata={
                "partition_key": partition_key,
                "p_date": p_date,
                "p_hour": p_hour,
                "today_count": today_count,
                "yesterday_count": yesterday_count,
                "change_rate": round(change_rate, 4),
                "change_percentage": f"+{round(change_rate * 100, 2)}%",
            },
            description=f"Volume increased or stable: +{round(change_rate * 100, 2)}%. Today: {today_count:,}, Yesterday: {yesterday_count:,}"
        )
    
    # Volume decreased - check if decrease is within acceptable range
    decrease_rate = (yesterday_count - today_count) / yesterday_count
    passed = decrease_rate <= MAX_DECREASE_RATE
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "partition_key": partition_key,
            "p_date": p_date,
            "p_hour": p_hour,
            "today_count": today_count,
            "yesterday_count": yesterday_count,
            "decrease_rate": round(decrease_rate, 4),
            "decrease_percentage": f"-{round(decrease_rate * 100, 2)}%",
            "max_allowed_decrease": f"{MAX_DECREASE_RATE * 100}%",
        },
        description=f"Volume decreased: -{round(decrease_rate * 100, 2)}% (threshold: -{MAX_DECREASE_RATE * 100}%). Today: {today_count:,}, Yesterday: {yesterday_count:,}"
    )