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

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


# Constants for validation
UUID_LENGTH = 36
MIN_ROW_COUNT = 0
MAX_NULL_RATE_USER_ID = 0.15
MAX_NULL_RATE_SONG_ID = 0.2
MAX_NULL_RATE_COUNTRY = 0.001
MAX_NULL_RATE_SESSION_ID = 0.01


@dg.asset_check(asset="web_audio_player_actions", 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_AUDIO_PLAYER_ACTIONS_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_audio_player_actions", blocking=False)
def check_user_id_null_rate(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check that USER_ID 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_ID IS NULL THEN 1 ELSE 0 END) as null_count
            FROM {WEB_AUDIO_PLAYER_ACTIONS_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={
            "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="web_audio_player_actions", blocking=False)
def check_song_id_null_rate(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check that SONG_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 SONG_ID IS NULL THEN 1 ELSE 0 END) as null_count
            FROM {WEB_AUDIO_PLAYER_ACTIONS_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_SONG_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_SONG_ID:.2%}",
        },
        description=f"SONG_ID null rate: {null_rate:.2%} (threshold: {MAX_NULL_RATE_SONG_ID:.2%})"
    )


@dg.asset_check(asset="web_audio_player_actions", 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_AUDIO_PLAYER_ACTIONS_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_audio_player_actions", 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_AUDIO_PLAYER_ACTIONS_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_audio_player_actions", blocking=False)
def check_distinct_users_and_songs(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    """Check distinct counts of users and songs 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_ID) as distinct_users,
                COUNT(DISTINCT SONG_ID) as distinct_songs
            FROM {WEB_AUDIO_PLAYER_ACTIONS_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_songs = result[2] if result and result[2] else 0
    
    # Basic sanity check: distinct users and songs should be > 0
    passed = distinct_users > 0 and distinct_songs > 0
    
    return dg.AssetCheckResult(
        passed=passed,
        metadata={
            "total_rows": total_rows,
            "distinct_users": distinct_users,
            "distinct_songs": distinct_songs,
            "avg_actions_per_user": round(total_rows / distinct_users, 2) if distinct_users > 0 else 0,
            "avg_actions_per_song": round(total_rows / distinct_songs, 2) if distinct_songs > 0 else 0,
        },
        description=f"Users: {distinct_users:,}, Songs: {distinct_songs:,}, Total actions: {total_rows:,}"
    )


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