"""Asset checks for rds_grouped_user_notification data quality."""

import dagster as dg
from dagster_snowflake import SnowflakeResource

from src.utils.snowflake.constants import Warehouse
from src.utils.snowflake.query import load_query
from src.utils.snowflake.partition_utils import parse_hourly_partition_key


@dg.asset_check(
    asset="rds_grouped_user_notification",
    description="Verify ID uniqueness",
    blocking=False
)
def check_grouped_user_notification_id_uniqueness(
    context: dg.AssetCheckExecutionContext,
    snowflake: SnowflakeResource
) -> dg.AssetCheckResult:
    """
    Check that all IDs in the partition are unique.
    
    Args:
        context: Asset check execution context
        snowflake: Snowflake resource
        
    Returns:
        AssetCheckResult indicating pass/fail status
    """
    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(
            load_query(
                "src/utils/snowflake/queries/use_warehouse.sql",
                params={"warehouse": Warehouse.SUNO_PROD_RDS_HOURLY_X_SMALL.value}
            )
        )
        cursor.execute(
            f"SELECT COUNT(*), COUNT(DISTINCT ID), COUNT(*) - COUNT(DISTINCT ID) "
            f"FROM SUNO_PROD.PROD.RDS_GROUPED_USER_NOTIFICATION "
            f"WHERE p_date = '{p_date}' AND p_hour = {p_hour}"
        )
        total, distinct, dupes = cursor.fetchone() or (0, 0, 0)
        
        return dg.AssetCheckResult(
            passed=dupes == 0,
            description=f"{'Unique' if dupes == 0 else f'{dupes} duplicates'}",
            metadata={
                "total_rows": total,
                "duplicates": dupes
            }
        )


@dg.asset_check(
    asset="rds_grouped_user_notification",
    description="Verify row count > 0",
    blocking=False
)
def check_grouped_user_notification_row_count(
    context: dg.AssetCheckExecutionContext,
    snowflake: SnowflakeResource
) -> dg.AssetCheckResult:
    """
    Check that the partition has at least one row.
    
    Args:
        context: Asset check execution context
        snowflake: Snowflake resource
        
    Returns:
        AssetCheckResult indicating pass/fail status
    """
    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(
            load_query(
                "src/utils/snowflake/queries/use_warehouse.sql",
                params={"warehouse": Warehouse.SUNO_PROD_RDS_HOURLY_X_SMALL.value}
            )
        )
        cursor.execute(
            f"SELECT COUNT(*) FROM SUNO_PROD.PROD.RDS_GROUPED_USER_NOTIFICATION "
            f"WHERE p_date = '{p_date}' AND p_hour = {p_hour}"
        )
        row_count = cursor.fetchone()[0]
        
        return dg.AssetCheckResult(
            passed=row_count > 0,
            description=f"Partition has {row_count} rows",
            metadata={"row_count": row_count}
        )


asset_checks = [
    check_grouped_user_notification_id_uniqueness,
    check_grouped_user_notification_row_count
]

__all__ = ["asset_checks"]