from typing import List

from pydantic import Field
from pyspark.sql import DataFrame
from pyspark.sql import functions as F

from sparkdq.core.base_check import BaseRowCheck
from sparkdq.core.base_config import BaseRowCheckConfig
from sparkdq.core.severity import Severity
from sparkdq.exceptions import MissingColumnError
from sparkdq.plugin.check_config_registry import register_check_config


class NotNullCheck(BaseRowCheck):
    """
    Row-level data quality check that flags rows where any of the specified columns is *not null*.

    This check appends a boolean result column to the input DataFrame.
    Each row will be marked as a failure (True) if any of the specified columns contain a value.
    This is useful in scenarios where columns are expected to remain empty (null) under normal conditions.

    If any of the target columns do not exist in the input DataFrame, a
    `MissingColumnError` is raised at runtime to signal a misconfiguration.

    Attributes:
        columns (List[str]): List of column names to validate for null values (i.e. should remain unset).
    """

    def __init__(self, check_id: str, columns: List[str], severity: Severity = Severity.CRITICAL):
        """
        Initialize a NotNullCheck instance.

        Args:
            check_id (str): Unique identifier for the check instance. Used as the
                name of the result column in the output DataFrame.
            columns (List[str]): List of column names that are expected to be null.
            severity (Severity, optional): Severity level of the check result.
                Defaults to Severity.CRITICAL.
        """
        super().__init__(check_id=check_id, severity=severity)
        self.columns = columns

    def validate(self, df: DataFrame) -> DataFrame:
        """
        Execute the not-null check on the given DataFrame.

        This method appends a new boolean column (named after `check_id`) that indicates
        for each row whether any of the specified columns is *not* null.

        Args:
            df (DataFrame): The input Spark DataFrame to validate.

        Returns:
            DataFrame: A new DataFrame with an additional boolean column where
                `True` indicates a non-null value (i.e. check failed), and `False` means valid.

        Raises:
            MissingColumnError: If any of the specified columns do not exist in the DataFrame.
        """
        for column in self.columns:
            if column not in df.columns:
                raise MissingColumnError(column, df.columns)

        # Build a Spark array where each element checks if the corresponding column is NOT NULL
        not_null_checks = F.array(*[F.col(c).isNotNull() for c in self.columns])

        # Reduce the array by OR-ing all not-null checks
        any_not_null_expr = F.aggregate(not_null_checks, F.lit(False), lambda acc, x: acc | x)

        return self.with_check_result_column(df, any_not_null_expr)


@register_check_config(check_name="not-null-check")
class NotNullCheckConfig(BaseRowCheckConfig):
    """
    Declarative configuration model for the NotNullCheck.

    Attributes:
        columns (List[str]): The names of the columns that should remain null.
    """

    check_class = NotNullCheck
    columns: List[str] = Field(..., description="The list of columns that must remain null")
