"""Asset checks for rds_onetime_free_usage data quality."""

import dagster as dg
from dagster_snowflake import SnowflakeResource

from src.utils.snowflake.constants import Warehouse
from src.utils.snowflake.logger import log_query
from src.utils.snowflake.query import load_query


@dg.asset_check(
    asset="rds_onetime_free_usage",
    description="Verify ID uniqueness in partition",
    blocking=False,
)
def check_rds_onetime_free_usage_id_uniqueness(
    context: dg.AssetCheckExecutionContext,
    snowflake: SnowflakeResource
) -> dg.AssetCheckResult:
    """Check that all IDs in the partition are unique."""
    logger = context.log
    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}))
        
        query = f"""
            SELECT COUNT(*) as total, COUNT(DISTINCT ID) as distinct_ids, COUNT(*) - COUNT(DISTINCT ID) as dupes
            FROM SUNO_PROD.PROD.RDS_ONETIME_FREE_USAGE
            WHERE p_date = '{p_date}' AND p_hour = {p_hour}
        """
        result = cursor.execute(query).fetchone()
        total, distinct_ids, dupes = result if result else (0, 0, 0)
    
    passed = dupes == 0
    return dg.AssetCheckResult(
        passed=passed,
        description=f"{'No duplicates' if passed else f'{dupes} duplicates found'}",
        metadata={"total_rows": total, "distinct_ids": distinct_ids, "duplicates": dupes}
    )


@dg.asset_check(
    asset="rds_onetime_free_usage",
    description="Verify row count > 0",
    blocking=False,
)
def check_rds_onetime_free_usage_row_count(
    context: dg.AssetCheckExecutionContext,
    snowflake: SnowflakeResource
) -> dg.AssetCheckResult:
    """Check that the partition has data (row count > 0)."""
    logger = context.log
    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}))
        
        query = f"""
            SELECT COUNT(*) FROM SUNO_PROD.PROD.RDS_ONETIME_FREE_USAGE
            WHERE p_date = '{p_date}' AND p_hour = {p_hour}
        """
        row_count = cursor.execute(query).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_rds_onetime_free_usage_id_uniqueness,
    check_rds_onetime_free_usage_row_count,
]

__all__ = ["asset_checks"]