import inspect
import logging
from enum import Enum
from dataclasses import dataclass, field
from abc import ABC, abstractmethod
import functools as ft
from typing import Any
from collections.abc import Callable
from datetime import datetime
from pyspark.sql import Column
import pyspark.sql.functions as F
from databricks.labs.dqx.utils import get_column_as_string

logger = logging.getLogger(__name__)


CHECK_FUNC_REGISTRY: dict[str, str] = {}


def register_rule(rule_type: str) -> Callable:
    def wrapper(func: Callable) -> Callable:
        CHECK_FUNC_REGISTRY[func.__name__] = rule_type
        return func

    return wrapper


class Criticality(Enum):
    """Enum class to represent criticality of the check."""

    WARN = "warn"
    ERROR = "error"


class DefaultColumnNames(Enum):
    """Enum class to represent columns in the dataframe that will be used for error and warning reporting."""

    ERRORS = "_errors"
    WARNINGS = "_warnings"


class ColumnArguments(Enum):
    """Enum class that is used as input parsing for custom column naming."""

    ERRORS = "errors"
    WARNINGS = "warnings"


@dataclass(frozen=True)
class ExtraParams:
    """Class to represent extra parameters for DQEngine."""

    column_names: dict[str, str] = field(default_factory=dict)
    run_time: datetime = datetime.now()
    user_metadata: dict[str, str] = field(default_factory=dict)


@dataclass(frozen=True)
class DQRule(ABC):
    """Represents a data quality rule that applies a quality check function to column(s) or
    column expression(s). This class includes the following attributes:
    * `check_func` - The function used to perform the quality check.
    * `name` (optional) - A custom name for the check; autogenerated if not provided.
    * `criticality` (optional) - Defines the severity level of the check:
        - `error`: Critical issues.
        - `warn`: Potential issues.
    * `filter` (optional) - A filter expression to apply the check only to rows meeting specific conditions.
    * `check_func_args` (optional) - Positional arguments for the check function (excluding `column`).
    * `check_func_kwargs` (optional) - Keyword arguments for the check function (excluding `column`).
    * `user_metadata` (optional) - User-defined key-value pairs added to metadata generated by the check.
    """

    check_func: Callable
    name: str = ""
    criticality: str = Criticality.ERROR.value
    filter: str | None = None
    check_func_args: list[Any] = field(default_factory=list)
    check_func_kwargs: dict[str, Any] = field(default_factory=dict)
    user_metadata: dict[str, str] | None = None

    def __post_init__(self):
        func_parameters = inspect.signature(self.check_func).parameters
        if "row_filter" in func_parameters:
            # pass filter if required by the check function (window type of checks)
            self.check_func_kwargs["row_filter"] = self.filter

        check = self._check

        object.__setattr__(self, "name", self.name if self.name else get_column_as_string(check, normalize=True))

    @ft.cached_property
    def check_criticality(self) -> str:
        """Criticality of the check.

        :return: string describing criticality - `warn` or `error`.
        :raises ValueError: if criticality is invalid.
        """
        criticality = self.criticality
        if criticality not in {Criticality.WARN.value, Criticality.ERROR.value}:
            raise ValueError(
                f"Invalid 'criticality' value: '{criticality}'. "
                f"Expected '{Criticality.WARN.value}' or '{Criticality.ERROR.value}'. "
                f"Check details: {self.name}"
            )
        return criticality

    @ft.cached_property
    def check_condition(self) -> Column:
        """Spark Column expression representing the check condition with filter.

        This expression returns a string value if the check evaluates to `true`,
        which serves as an error or warning message. If the check evaluates to `false`,
        it returns `null`. If a filter condition is provided, the check is applied
        only to rows that satisfy the condition.

        :return: A Spark Column object representing the check condition.
        """
        # if filter is provided, apply the filter to the check
        filter_col = F.expr(self.filter) if self.filter else F.lit(True)
        check = self._check
        return F.when(check.isNotNull(), F.when(filter_col, check)).otherwise(check)

    @ft.cached_property
    def columns_as_string_expr(self) -> Column:
        """Spark Column expression representing the column(s) as a string (not normalized).
        :return: A Spark Column object representing the column(s) as a string (not normalized).
        """
        return F.lit(None).cast("array<string>")

    @abstractmethod
    @ft.cached_property
    def _check(self) -> Column:
        """Spark Column expression representing the check condition.
        :return: A Spark Column object representing the check condition.
        """


@dataclass(frozen=True)
class DQRowSingleColRule(DQRule):
    """Represents a row-level data quality rule that applies a quality check function to a column or column expression.
    Works with check functions that take a single column or no column as input.
    This class extends DQRule and includes the following attributes in addition:
    * `column` - A single column to which the check function is applied."""

    column: str | Column | None = None

    def __post_init__(self):
        rule_type = CHECK_FUNC_REGISTRY.get(self.check_func.__name__)
        if rule_type and rule_type not in ("single_column", "no_column"):
            raise ValueError(
                f"Function '{self.check_func.__name__}' is not a single-column rule. Use DQRowMultiColRule instead."
            )

        super().__post_init__()

    @ft.cached_property
    def columns_as_string_expr(self) -> Column:
        """Spark Column expression representing the column(s) as a string (not normalized).
        :return: A Spark Column object representing the column(s) as a string (not normalized).
        """
        if self.column is not None:
            return F.array(F.lit(get_column_as_string(self.column)))
        return super().columns_as_string_expr

    @ft.cached_property
    def _check(self) -> Column:
        """Spark Column expression representing the check condition.
        :return: A Spark Column object representing the check condition.
        """
        args = [self.column] if self.column is not None else []
        args.extend(self.check_func_args)
        return self.check_func(*args, **self.check_func_kwargs)


@dataclass(frozen=True)
class DQRowMultiColRule(DQRule):
    """Represents a row-level data quality rule that applies a quality check function to multiple columns or
    column expressions. Works only with check functions that take list of columns as input.
    This class extends DQRule and includes the following attributes in addition:
    * `columns` - Columns to which the check function is applied."""

    columns: list[str | Column] | None = None

    def __post_init__(self):
        rule_type = CHECK_FUNC_REGISTRY.get(self.check_func.__name__)
        if rule_type and rule_type != "multi_column":
            raise ValueError(
                f"Function '{self.check_func.__name__}' is not a multi-column rule. Use DQRowSingleColRule instead."
            )

        super().__post_init__()

    @ft.cached_property
    def columns_as_string_expr(self) -> Column:
        """Spark Column expression representing the column(s) as a string (not normalized).
        :return: A Spark Column object representing the column(s) as a string (not normalized).
        """
        if self.columns is not None:
            return F.array(*[F.lit(get_column_as_string(column)) for column in self.columns])
        return super().columns_as_string_expr

    @ft.cached_property
    def _check(self) -> Column:
        """Spark Column expression representing the check condition.
        :return: A Spark Column object representing the check condition.
        """
        args = [self.columns] if self.columns is not None else []
        args.extend(self.check_func_args)
        return self.check_func(*args, **self.check_func_kwargs)


def DQRowRule(*, column: str | Column | None = None, columns: list[str | Column] | None = None, **kwargs) -> DQRule:
    """Factory function to create a row-level data quality rule (check).
    :param column: A single column to which the check function is applied.
    :param columns: A list of columns to which the check function is applied.
    :param kwargs: Additional keyword arguments to pass to the rule.
    :return: An instance of DQRowSingleColRule or DQRowMultiColRule.
    :raises ValueError: If both `column` and `columns` are provided.
    """
    if column is not None and columns is not None:
        raise ValueError("Both 'column' and 'columns' cannot be provided at the same time.")

    if columns is not None:
        return DQRowMultiColRule(columns=columns, **kwargs)
    return DQRowSingleColRule(column=column, **kwargs)


@dataclass(frozen=True)
class DQRowRuleForEachCol:
    """Represents a row-level data quality rule set that applies a quality check function
    repeatedly on each specified column of the provided list of columns.
    This class includes the following attributes:
    * `columns` - A list of column names or expressions to which the check function should be applied.
    * `check_func` - The function used to perform the quality check.
    * `name` (optional) - A custom name for the check; autogenerated if not provided.
    * `criticality` - The severity level of the check:
        - `'warn'` for potential issues.
        - `'error'` for critical issues.
    * `filter` (optional) - A filter expression to apply the check only to rows meeting specific conditions.
    * `check_func_args` (optional) - Positional arguments for the check function (excluding column names).
    * `check_func_kwargs` (optional) - Keyword arguments for the check function (excluding column names).
    * `user_metadata` (optional) - User-defined key-value pairs added to metadata generated by the check.
    """

    columns: list[str | Column | list[str | Column]]
    check_func: Callable
    name: str = ""
    criticality: str = Criticality.ERROR.value
    filter: str | None = None
    check_func_args: list[Any] = field(default_factory=list)
    check_func_kwargs: dict[str, Any] = field(default_factory=dict)
    user_metadata: dict[str, str] | None = None

    def get_rules(self) -> list[DQRule]:
        """Build a list of rules for a set of columns.

        :return: list of dq rules
        """
        rules: list[DQRule] = []
        for col_set in self.columns:
            if isinstance(col_set, list):
                multi_col_rule = DQRowMultiColRule(
                    columns=col_set,
                    name=self.name,
                    criticality=self.criticality,
                    check_func=self.check_func,
                    check_func_args=self.check_func_args,
                    check_func_kwargs=self.check_func_kwargs,
                    filter=self.filter,
                    user_metadata=self.user_metadata,
                )
                rules.append(multi_col_rule)
            else:
                col_rule = DQRowSingleColRule(
                    column=col_set,
                    name=self.name,
                    criticality=self.criticality,
                    check_func=self.check_func,
                    check_func_args=self.check_func_args,
                    check_func_kwargs=self.check_func_kwargs,
                    filter=self.filter,
                    user_metadata=self.user_metadata,
                )
                rules.append(col_rule)
        return rules


@dataclass(frozen=True)
class ChecksValidationStatus:
    """Class to represent the validation status."""

    _errors: list[str] = field(default_factory=list)

    def add_error(self, error: str):
        """Add an error to the validation status."""
        self._errors.append(error)

    def add_errors(self, errors: list[str]):
        """Add an error to the validation status."""
        self._errors.extend(errors)

    @property
    def has_errors(self) -> bool:
        """Check if there are any errors in the validation status."""
        return bool(self._errors)

    @property
    def errors(self) -> list[str]:
        """Get the list of errors in the validation status."""
        return self._errors

    def to_string(self) -> str:
        """Convert the validation status to a string."""
        if self.has_errors:
            return "\n".join(self._errors)
        return "No errors found"

    def __str__(self) -> str:
        """String representation of the ValidationStatus class."""
        return self.to_string()
