"""Asset checks for rds_hook_reaction data quality."""

from datetime import timedelta

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_hook_reaction",
    description="Verify that all IDs in the partition are unique (no duplicates)",
    blocking=False,
)
def check_hook_reaction_id_uniqueness(
    context: dg.AssetCheckExecutionContext,
    snowflake: SnowflakeResource
) -> dg.AssetCheckResult:
    """
    Check that all IDs in the current partition are unique.
    
    This check:
    1. Counts total rows in the partition
    2. Counts distinct IDs in the partition
    3. Identifies any duplicate IDs if found
    4. Fails if duplicates exist, passes otherwise
    """
    logger = context.log
    
    # Get partition information - for asset checks, use op_execution_context
    partition_key = context.op_execution_context.partition_key
    logger.info(f"Checking ID uniqueness for partition: {partition_key}")
    
    # Parse partition key (format: YYYY-MM-DD-HH:00)
    # Example: "2025-10-03-14:00"
    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])
    
    logger.info(f"Checking partition: p_date={p_date}, p_hour={p_hour}")
    
    with snowflake.get_connection() as conn:
        cursor = conn.cursor()
        
        # Use warehouse
        warehouse_query = load_query(
            "src/utils/snowflake/queries/use_warehouse.sql",
            params={"warehouse": Warehouse.SUNO_PROD_RDS_HOURLY_X_SMALL.value}
        )
        cursor.execute(warehouse_query)
        
        # Query 1: Check for duplicate IDs
        duplicate_check_query = f"""
            SELECT 
                COUNT(*) as total_rows,
                COUNT(DISTINCT ID) as distinct_ids,
                COUNT(*) - COUNT(DISTINCT ID) as duplicate_count
            FROM SUNO_PROD.PROD.RDS_HOOK_REACTION
            WHERE p_date = '{p_date}'
              AND p_hour = {p_hour}
        """
        
        log_query(logger, duplicate_check_query)
        cursor.execute(duplicate_check_query)
        result = cursor.fetchone()
        
        if not result:
            return dg.AssetCheckResult(
                passed=False,
                description="No data found in partition",
                metadata={
                    "partition_key": partition_key,
                    "p_date": p_date,
                    "p_hour": p_hour,
                }
            )
        
        total_rows, distinct_ids, duplicate_count = result
        
        logger.info(f"Total rows: {total_rows}, Distinct IDs: {distinct_ids}, Duplicates: {duplicate_count}")
        
        # If there are duplicates, get details
        if duplicate_count > 0:
            return dg.AssetCheckResult(
                passed=False,
                description=f"Found {duplicate_count} duplicate IDs in partition {partition_key}",
                metadata={
                    "partition_key": partition_key,
                    "p_date": p_date,
                    "p_hour": p_hour,
                    "total_rows": total_rows,
                    "distinct_ids": distinct_ids,
                    "duplicate_count": duplicate_count,
                }
            )
        
        # No duplicates - check passes
        return dg.AssetCheckResult(
            passed=True,
            description=f"All {total_rows} IDs in partition {partition_key} are unique",
            metadata={
                "partition_key": partition_key,
                "p_date": p_date,
                "p_hour": p_hour,
                "total_rows": total_rows,
                "distinct_ids": distinct_ids,
                "duplicate_count": 0,
            }
        )


@dg.asset_check(
    asset="rds_hook_reaction",
    description="Verify that the partition has at least a minimum row count",
    blocking=False,
)
def check_hook_reaction_row_count(
    context: dg.AssetCheckExecutionContext,
    snowflake: SnowflakeResource
) -> dg.AssetCheckResult:
    """
    Check that the partition has data (row count > 0).
    
    This check validates that the partition was processed successfully
    and contains at least one row.
    """
    logger = context.log
    
    # Get partition information - for asset checks, use op_execution_context
    partition_key = context.op_execution_context.partition_key
    logger.info(f"Checking row count for partition: {partition_key}")
    
    # Parse partition key (format: YYYY-MM-DD-HH:00)
    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])
    
    logger.info(f"Checking partition: p_date={p_date}, p_hour={p_hour}")
    
    with snowflake.get_connection() as conn:
        cursor = conn.cursor()
        
        # Use warehouse
        warehouse_query = load_query(
            "src/utils/snowflake/queries/use_warehouse.sql",
            params={"warehouse": Warehouse.SUNO_PROD_RDS_HOURLY_X_SMALL.value}
        )
        cursor.execute(warehouse_query)
        
        # Check row count
        row_count_query = f"""
            SELECT COUNT(*) as row_count
            FROM SUNO_PROD.PROD.RDS_HOOK_REACTION
            WHERE p_date = '{p_date}'
              AND p_hour = {p_hour}
        """
        
        log_query(logger, row_count_query)
        cursor.execute(row_count_query)
        result = cursor.fetchone()
        
        row_count = result[0] if result else 0
        
        logger.info(f"Row count: {row_count}")
        
        # Pass if row count > 0 (require at least 1 row)
        passed = row_count > 0
        
        return dg.AssetCheckResult(
            passed=passed,
            description=f"Partition has {row_count} rows",
            metadata={
                "partition_key": partition_key,
                "p_date": p_date,
                "p_hour": p_hour,
                "row_count": row_count,
            }
        )