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

import dagster as dg
from dagster_snowflake import SnowflakeResource
from src.assets.snowflake.raw_table.frontend.web_hook_event.assets import WEB_HOOK_EVENT_TABLE_NAME
from src.utils.snowflake.partition_utils import parse_hourly_partition_key


# Constants for validation
UUID_LENGTH = 36
LONG_UUID_LENGTH = 64
MIN_ROW_COUNT = 0
MAX_NULL_RATE_USER_UID = 0.15
MAX_NULL_RATE_COUNTRY = 0.001
MAX_NULL_RATE_ANONYMOUS_ID = 0.001
MAX_NULL_RATE_SESSION_ID = 0.01
MAX_INVALID_LENGTH_RATE = 0.01


@dg.asset_check(asset="web_hook_event", blocking=False)
def check_minimum_row_count(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Ensure minimum row count threshold is met (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()
        query = f"""
            SELECT COUNT(*) as row_count
            FROM {WEB_HOOK_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={
            "row_count": row_count,
            "min_threshold": MIN_ROW_COUNT,
            "partition_key": partition_key,
        },
        description=f"Row count: {row_count:,} (threshold: {MIN_ROW_COUNT:,})"
    )


@dg.asset_check(asset="web_hook_event", blocking=False)
def check_user_uid_null_rate(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check that USER_UID null rate is below 15%."""
    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()
        query = f"""
            SELECT 
                COUNT(*) as total_count,
                SUM(CASE WHEN USER_UID IS NULL THEN 1 ELSE 0 END) as null_count
            FROM {WEB_HOOK_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_UID
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "null_count": null_count,
            "total_count": total_count,
            "null_rate": f"{null_rate:.2%}",
            "threshold": f"{MAX_NULL_RATE_USER_UID:.2%}",
        },
        description=f"USER_UID null rate: {null_rate:.2%} (threshold: {MAX_NULL_RATE_USER_UID:.2%})"
    )


@dg.asset_check(asset="web_hook_event", blocking=False)
def check_country_null_rate(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check that COUNTRY null rate is below 0.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()
        query = f"""
            SELECT 
                COUNT(*) as total_count,
                SUM(CASE WHEN COUNTRY IS NULL OR COUNTRY = '' THEN 1 ELSE 0 END) as null_count
            FROM {WEB_HOOK_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_COUNTRY
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "null_count": null_count,
            "total_count": total_count,
            "null_rate": f"{null_rate:.4%}",
            "threshold": f"{MAX_NULL_RATE_COUNTRY:.2%}",
        },
        description=f"COUNTRY null rate: {null_rate:.4%}"
    )


@dg.asset_check(asset="web_hook_event", blocking=False)
def check_session_id_null_rate(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check that SESSION_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()
        query = f"""
            SELECT 
                COUNT(*) as total_count,
                SUM(CASE WHEN SESSION_ID IS NULL THEN 1 ELSE 0 END) as null_count
            FROM {WEB_HOOK_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={
            "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%}"
    )


@dg.asset_check(asset="web_hook_event", blocking=False)
def check_anonymous_id_quality(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check ANONYMOUS_ID null rate and length."""
    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()
        query = f"""
            SELECT 
                COUNT(*) as total_count,
                SUM(CASE WHEN ANONYMOUS_ID IS NULL THEN 1 ELSE 0 END) as null_count,
                SUM(CASE 
                    WHEN ANONYMOUS_ID IS NOT NULL AND LENGTH(ANONYMOUS_ID) != {UUID_LENGTH} 
                    THEN 1 ELSE 0 
                END) as invalid_length_count
            FROM {WEB_HOOK_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
        invalid_count = result[2] if result and result[2] else 0
    
    null_rate = null_count / total_count if total_count > 0 else 0
    invalid_rate = invalid_count / total_count if total_count > 0 else 0
    
    passed = null_rate <= MAX_NULL_RATE_ANONYMOUS_ID and invalid_rate <= MAX_INVALID_LENGTH_RATE
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "null_rate": f"{null_rate:.4%}",
            "invalid_length_rate": f"{invalid_rate:.2%}",
            "expected_length": UUID_LENGTH,
        },
        description=f"ANONYMOUS_ID quality check - null: {null_rate:.4%}, invalid length: {invalid_rate:.2%}"
    )


@dg.asset_check(asset="web_hook_event", blocking=False)
def check_distinct_users_and_sessions(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check distinct counts of users and sessions for data volume validation."""
    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()
        query = f"""
            SELECT 
                COUNT(*) as total_rows,
                COUNT(DISTINCT USER_UID) as distinct_users,
                COUNT(DISTINCT SESSION_ID) as distinct_sessions
            FROM {WEB_HOOK_EVENT_TABLE_NAME}
            WHERE p_date = '{p_date}' and p_hour = {p_hour}
        """
        result = cursor.execute(query).fetchone()
        total_rows = result[0] if result else 0
        distinct_users = result[1] if result and result[1] else 0
        distinct_sessions = result[2] if result and result[2] else 0
    
    # Basic sanity check: distinct users and sessions should be > 0
    passed = distinct_users > 0 and distinct_sessions > 0
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "total_rows": total_rows,
            "distinct_users": distinct_users,
            "distinct_sessions": distinct_sessions,
            "avg_events_per_user": round(total_rows / distinct_users, 2) if distinct_users > 0 else 0,
            "avg_events_per_session": round(total_rows / distinct_sessions, 2) if distinct_sessions > 0 else 0,
        },
        description=f"Users: {distinct_users:,}, Sessions: {distinct_sessions:,}, Total rows: {total_rows:,}"
    )


@dg.asset_check(asset="web_hook_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()
        
        # 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 {WEB_HOOK_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_uid_null_rate,
    check_country_null_rate,
    check_session_id_null_rate,
    check_anonymous_id_quality,
    check_distinct_users_and_sessions,
    check_day_over_day_volume,
]

__all__ = ["asset_checks"]