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_period_credit_usage", description="Verify ID uniqueness in partition", blocking=False)
def check_period_credit_usage_id_uniqueness(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    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_MEDUIM.value}))
        cursor.execute(f"SELECT COUNT(*), COUNT(DISTINCT ID), COUNT(*) - COUNT(DISTINCT ID) FROM SUNO_PROD.PROD.RDS_PERIOD_CREDIT_USAGE WHERE p_date = '{p_date}' AND p_hour = {p_hour}")
        total, distinct, dupes = cursor.fetchone() or (0, 0, 0)
        
    passed = dupes == 0
    return dg.AssetCheckResult(passed=passed, description=f"{'All IDs unique' if passed else f'{dupes} duplicate IDs found'}", metadata={"total_rows": total, "distinct_ids": distinct, "duplicates": dupes})

@dg.asset_check(asset="rds_period_credit_usage", description="Verify row count > 0", blocking=False)
def check_period_credit_usage_row_count(context: dg.AssetCheckExecutionContext, snowflake: SnowflakeResource) -> dg.AssetCheckResult:
    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_MEDUIM.value}))
        cursor.execute(f"SELECT COUNT(*) FROM SUNO_PROD.PROD.RDS_PERIOD_CREDIT_USAGE WHERE p_date = '{p_date}' AND p_hour = {p_hour}")
        row_count = cursor.fetchone()[0]
    
    passed = row_count > 0
    return dg.AssetCheckResult(passed=passed, description=f"Partition has {row_count} rows (expected > 0)", metadata={"row_count": row_count})

asset_checks = [check_period_credit_usage_id_uniqueness, check_period_credit_usage_row_count]
__all__ = ["asset_checks"]