"""Asset checks for rds_discount_offer 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

@dg.asset_check(asset="rds_discount_offer", description="Verify ID uniqueness", blocking=False)
def check_discount_offer_id_uniqueness(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    partition_key = context.op_execution_context.partition_key
    partition_parts = partition_key.split("-")
    p_date = f"{partition_parts[0]}-{partition_parts[1]}-{partition_parts[2]}"
    p_hour = int(partition_parts[3].split(":")[0])
    
    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(*) as total_rows, COUNT(DISTINCT ID) as distinct_ids, COUNT(*) - COUNT(DISTINCT ID) as duplicate_count FROM SUNO_PROD.PROD.RDS_DISCOUNT_OFFER WHERE p_date = '{p_date}' AND p_hour = {p_hour}")
        result = cursor.fetchone()
        
        if not result:
            return dg.AssetCheckResult(passed=False, description="No data found")
        
        total_rows, distinct_ids, duplicate_count = result
        passed = duplicate_count == 0
        
        return dg.AssetCheckResult(
            passed=passed,
            description=f"{'All IDs unique' if passed else f'Found {duplicate_count} duplicates'}",
            metadata={"total_rows": total_rows, "distinct_ids": distinct_ids, "duplicate_count": duplicate_count}
        )

asset_checks = [check_discount_offer_id_uniqueness]
__all__ = ["asset_checks"]