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

import dagster as dg
from dagster_snowflake import SnowflakeResource
from src.assets.snowflake.raw_table.frontend.app_event.assets import APP_EVENT_TABLE_NAME, WAREHOUSE
from src.utils.snowflake.partition_utils import parse_hourly_partition_key


# Constants for validation
MIN_ROW_COUNT = 0  # Allow partitions with 0 rows
MAX_NULL_RATE_USER_ID = 0.50  # Allow up to 50% null USER_ID (anonymous users)
MAX_NULL_RATE_ANONYMOUS_ID = 0.05  # Most events should have anonymous_id
MAX_NULL_RATE_SESSION_ID = 0.3  # Most events should have session
MAX_VOLUME_CHANGE_RATE = 0.10  # Maximum 10% change from previous day


@dg.asset_check(asset="app_event", blocking=False)
def check_minimum_row_count(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Ensure minimum row count threshold is met (at least 0 rows)."""
    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 {APP_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="app_event", blocking=False)
def check_user_id_null_rate(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check that USER_ID null rate is below 50% (many users are anonymous)."""
    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 total_count,
                SUM(CASE WHEN USER_ID IS NULL OR USER_ID = '' THEN 1 ELSE 0 END) as null_count
            FROM {APP_EVENT_TABLE_NAME}
            WHERE p_date = '{p_date}' and p_hour = {p_hour}
        """
        result = cursor.execute(query).fetchone()
        total_count = result[0] if result else 0
        null_count = result[1] if result and result[1] else 0
    
    null_rate = null_count / total_count if total_count > 0 else 0
    passed = null_rate <= MAX_NULL_RATE_USER_ID
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "partition_key": partition_key,
            "p_date": p_date,
            "p_hour": p_hour,
            "null_count": null_count,
            "total_count": total_count,
            "null_rate": f"{null_rate:.2%}",
            "threshold": f"{MAX_NULL_RATE_USER_ID:.2%}",
        },
        description=f"USER_ID null rate: {null_rate:.2%} (threshold: {MAX_NULL_RATE_USER_ID:.2%})"
    )


@dg.asset_check(asset="app_event", blocking=False)
def check_anonymous_id_null_rate(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check that ANONYMOUS_ID null rate is below 1%."""
    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 total_count,
                SUM(CASE WHEN ANONYMOUS_ID IS NULL OR ANONYMOUS_ID = '' THEN 1 ELSE 0 END) as null_count
            FROM {APP_EVENT_TABLE_NAME}
            WHERE p_date = '{p_date}' and p_hour = {p_hour}
        """
        result = cursor.execute(query).fetchone()
        total_count = result[0] if result else 0
        null_count = result[1] if result and result[1] else 0
    
    null_rate = null_count / total_count if total_count > 0 else 0
    passed = null_rate <= MAX_NULL_RATE_ANONYMOUS_ID
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "partition_key": partition_key,
            "p_date": p_date,
            "p_hour": p_hour,
            "null_count": null_count,
            "total_count": total_count,
            "null_rate": f"{null_rate:.2%}",
            "threshold": f"{MAX_NULL_RATE_ANONYMOUS_ID:.2%}",
        },
        description=f"ANONYMOUS_ID null rate: {null_rate:.2%} (threshold: {MAX_NULL_RATE_ANONYMOUS_ID:.2%})"
    )


@dg.asset_check(asset="app_event", blocking=False)
def check_session_id_null_rate(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check that SESSION_ID null rate is below 5%."""
    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 total_count,
                SUM(CASE WHEN SESSION_ID IS NULL OR SESSION_ID = '' THEN 1 ELSE 0 END) as null_count
            FROM {APP_EVENT_TABLE_NAME}
            WHERE p_date = '{p_date}' and p_hour = {p_hour}
        """
        result = cursor.execute(query).fetchone()
        total_count = result[0] if result else 0
        null_count = result[1] if result and result[1] else 0
    
    null_rate = null_count / total_count if total_count > 0 else 0
    passed = null_rate <= MAX_NULL_RATE_SESSION_ID
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "partition_key": partition_key,
            "p_date": p_date,
            "p_hour": p_hour,
            "null_count": null_count,
            "total_count": total_count,
            "null_rate": f"{null_rate:.2%}",
            "threshold": f"{MAX_NULL_RATE_SESSION_ID:.2%}",
        },
        description=f"SESSION_ID null rate: {null_rate:.2%} (threshold: {MAX_NULL_RATE_SESSION_ID:.2%})"
    )


@dg.asset_check(asset="app_event", blocking=False)
def check_distinct_users_and_sessions(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check distinct user and session counts."""
    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 total_count,
                COUNT(DISTINCT USER_ID) as distinct_users,
                COUNT(DISTINCT ANONYMOUS_ID) as distinct_anonymous,
                COUNT(DISTINCT SESSION_ID) as distinct_sessions
            FROM {APP_EVENT_TABLE_NAME}
            WHERE p_date = '{p_date}' and p_hour = {p_hour}
        """
        result = cursor.execute(query).fetchone()
        total_count = result[0] if result else 0
        distinct_users = result[1] if result else 0
        distinct_anonymous = result[2] if result else 0
        distinct_sessions = result[3] if result else 0
    
    # Calculate average events per user/session
    avg_events_per_user = total_count / distinct_users if distinct_users > 0 else 0
    avg_events_per_session = total_count / distinct_sessions if distinct_sessions > 0 else 0
    
    # Pass if we have distinct values (or total_count is 0)
    passed = total_count == 0 or (distinct_users > 0 or distinct_anonymous > 0) and distinct_sessions > 0
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "partition_key": partition_key,
            "p_date": p_date,
            "p_hour": p_hour,
            "total_count": total_count,
            "distinct_users": distinct_users,
            "distinct_anonymous": distinct_anonymous,
            "distinct_sessions": distinct_sessions,
            "avg_events_per_user": round(avg_events_per_user, 2),
            "avg_events_per_session": round(avg_events_per_session, 2),
        },
        description=f"Users: {distinct_users:,}, Anonymous: {distinct_anonymous:,}, Sessions: {distinct_sessions:,}"
    )


@dg.asset_check(asset="app_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 {APP_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:,}"
    )


# Export all asset checks
asset_checks = [
    check_minimum_row_count,
    check_user_id_null_rate,
    check_anonymous_id_null_rate,
    check_session_id_null_rate,
    check_distinct_users_and_sessions,
    check_day_over_day_volume,
]

__all__ = ["asset_checks"]