import logging
import warnings
from enum import Enum
from dataclasses import dataclass, field
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__)


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 DQColRule:
    """Represents a row-level data quality rule that applies a quality check function to a column
    or an SQL expression. This class includes the following attributes:
    * `check_func` - The function used to perform the quality check.
    * `col_name` - The name of the column to which the check function is applied.
    * `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 `col_name`).
    * `check_func_kwargs` (optional) - Keyword arguments for the check function (excluding `col_name`).
    """

    check_func: Callable
    col_name: str | Column | None = None
    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)

    def __post_init__(self):
        # validates correct args and kwargs are passed
        check = self._check

        # take the name from the alias of the column expression if not provided
        object.__setattr__(
            self, "name", self.name if self.name else "col_" + get_column_as_string(check, normalize=True)
        )

    @ft.cached_property
    def col_name_as_string_expr(self) -> Column:
        """Spark Column expression representing the column name as a string (not normalized).

        :return: A Spark Column object representing the column name as a string (not normalized).
        """
        if self.col_name is not None:
            return F.lit(get_column_as_string(self.col_name))
        return F.lit(self.col_name).cast("string")

    @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 _check(self) -> Column:
        """Spark Column expression representing the check condition.

        :return: A Spark Column object representing the check condition.
        """
        args = [self.col_name] if self.col_name is not None else []
        args.extend(self.check_func_args)
        return self.check_func(*args, **self.check_func_kwargs)


@dataclass(frozen=True)
class DQRule(DQColRule):
    """Represents a row-level data quality rule that applies a quality check function to a column.
    DQRule is deprecated and will be removed in a future version. Please use DQColRule instead.
    """

    def __init__(self, *args, **kwargs):
        depreciation_msg = "DQRule is deprecated and will be removed in a future version. Please use DQColRule instead."
        warnings.warn(depreciation_msg, DeprecationWarning, stacklevel=2)
        logger.warning(depreciation_msg)
        super().__init__(*args, **kwargs)


@dataclass(frozen=True)
class DQColSetRule:
    """Represents a row-level data quality rule set that applies a quality check function to multiple columns.
    This class includes the following attributes:
    * `columns` - A list of column names to which the check function should be applied.
    * `criticality` - The severity level of the check:
        - `'warn'` for potential issues.
        - `'error'` for critical issues.
    * `check_func` - The function used to perform the quality check.
    * `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).
    """

    columns: 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)

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

        :return: list of dq rules
        """
        rules = []
        for col_name in self.columns:
            rule = DQColRule(
                col_name=col_name,
                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,
            )
            rules.append(rule)
        return rules


@dataclass(frozen=True)
class DQRuleColSet(DQColSetRule):
    """Represents a row-level data quality rule that applies a quality check function to multiple columns.
    DQRuleColSet is deprecated and will be removed in a future version. Please use DQColSetRule instead.
    """

    def __init__(self, *args, **kwargs):
        depreciation_msg = (
            "DQRuleColSet is deprecated and will be removed in a future version. Please use DQRuleColSet instead."
        )
        warnings.warn(depreciation_msg, DeprecationWarning, stacklevel=2)
        logger.warning(depreciation_msg)
        super().__init__(*args, **kwargs)


@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()
