import abc
from functools import cached_property
from typing import Any, final
from pyspark.sql import DataFrame
from databricks.labs.dqx.rule import ChecksValidationStatus, DQRule
from databricks.sdk import WorkspaceClient
from databricks.labs.dqx.__about__ import __version__


class DQEngineBase(abc.ABC):
    def __init__(self, workspace_client: WorkspaceClient):
        self._workspace_client = workspace_client

    @cached_property
    def ws(self) -> WorkspaceClient:
        """
        Cached property to verify and return the workspace client.
        """
        return self._verify_workspace_client(self._workspace_client)

    @staticmethod
    @final
    def _verify_workspace_client(ws: WorkspaceClient) -> WorkspaceClient:
        """
        Verifies the Databricks workspace client configuration.
        """
        # Using reflection to set right value for _product_info as dqx for telemetry
        product_info = getattr(ws.config, '_product_info')
        if product_info[0] != "dqx":
            setattr(ws.config, '_product_info', ('dqx', __version__))

        # make sure Databricks workspace is accessible
        ws.current_user.me()
        return ws


class DQEngineCoreBase(DQEngineBase):

    @abc.abstractmethod
    def apply_checks(self, df: DataFrame, checks: list[DQRule]) -> DataFrame:
        """Applies data quality checks to a given dataframe.

        :param df: dataframe to check
        :param checks: list of checks to apply to the dataframe. Each check is an instance of DQRule class.
        :return: dataframe with errors and warning reporting columns
        """

    @abc.abstractmethod
    def apply_checks_and_split(self, df: DataFrame, checks: list[DQRule]) -> tuple[DataFrame, DataFrame]:
        """Applies data quality checks to a given dataframe and split it into two ("good" and "bad"),
        according to the data quality checks.

        :param df: dataframe to check
        :param checks: list of checks to apply to the dataframe. Each check is an instance of DQRule class.
        :return: two dataframes - "good" which includes warning rows but no reporting columns, and "data" having
        error and warning rows and corresponding reporting columns
        """

    @abc.abstractmethod
    def apply_checks_by_metadata_and_split(
        self, df: DataFrame, checks: list[dict], custom_check_functions: dict[str, Any] | None = None
    ) -> tuple[DataFrame, DataFrame]:
        """Wrapper around `apply_checks_and_split` for use in the metadata-driven pipelines. The main difference
        is how the checks are specified - instead of using functions directly, they are described as function name plus
        arguments.

        :param df: dataframe to check
        :param checks: list of dictionaries describing checks. Each check is a dictionary consisting of following fields:
        * `check` - Column expression to evaluate. This expression should return string value if it's evaluated to true -
        it will be used as an error/warning message, or `null` if it's evaluated to `false`
        * `name` - name that will be given to a resulting column. Autogenerated if not provided
        * `criticality` (optional) - possible values are `error` (data going only into "bad" dataframe),
        and `warn` (data is going into both dataframes)
        :param custom_check_functions: dictionary with custom check functions (eg. ``globals()`` of the calling module).
        If not specified, then only built-in functions are used for the checks.
        :return: two dataframes - "good" which includes warning rows but no reporting columns, and "bad" having
        error and warning rows and corresponding reporting columns
        """

    @abc.abstractmethod
    def apply_checks_by_metadata(
        self, df: DataFrame, checks: list[dict], custom_check_functions: dict[str, Any] | None = None
    ) -> DataFrame:
        """Wrapper around `apply_checks` for use in the metadata-driven pipelines. The main difference
        is how the checks are specified - instead of using functions directly, they are described as function name plus
        arguments.

        :param df: dataframe to check
        :param checks: list of dictionaries describing checks. Each check is a dictionary consisting of following fields:
        * `check` - Column expression to evaluate. This expression should return string value if it's evaluated to true -
        it will be used as an error/warning message, or `null` if it's evaluated to `false`
        * `name` - name that will be given to a resulting column. Autogenerated if not provided
        * `criticality` (optional) - possible values are `error` (data going only into "bad" dataframe),
        and `warn` (data is going into both dataframes)
        :param custom_check_functions: dictionary with custom check functions (eg. ``globals()`` of calling module).
        If not specified, then only built-in functions are used for the checks.
        :return: dataframe with errors and warning reporting columns
        """

    @staticmethod
    @abc.abstractmethod
    def validate_checks(
        checks: list[dict], custom_check_functions: dict[str, Any] | None = None
    ) -> ChecksValidationStatus:
        """
        Validate the input dict to ensure they conform to expected structure and types.

        Each check can be a dictionary. The function validates
        the presence of required keys, the existence and callability of functions, and the types
        of arguments passed to these functions.

        :param checks: List of checks to apply to the dataframe. Each check should be a dictionary.
        :param custom_check_functions: Optional dictionary with custom check functions.

        :return ValidationStatus: The validation status.
        """

    @abc.abstractmethod
    def get_invalid(self, df: DataFrame) -> DataFrame:
        """
        Get records that violate data quality checks (records with warnings and errors).
        @param df: input DataFrame.
        @return: dataframe with error and warning rows and corresponding reporting columns.
        """

    @abc.abstractmethod
    def get_valid(self, df: DataFrame) -> DataFrame:
        """
        Get records that don't violate data quality checks (records with warnings but no errors).
        @param df: input DataFrame.
        @return: dataframe with warning rows but no reporting columns.
        """

    @staticmethod
    @abc.abstractmethod
    def load_checks_from_local_file(filepath: str) -> list[dict]:
        """
        Load checks (dq rules) from a file (json or yaml) in the local file system.
        This does not require installation of DQX in the workspace.
        The returning checks can be used as input for `apply_checks_by_metadata` function.

        :param filepath: path to a file containing the checks.
        :return: list of dq rules
        """

    @staticmethod
    @abc.abstractmethod
    def save_checks_in_local_file(checks: list[dict], filepath: str):
        """
        Save checks (dq rules) to yaml file in the local file system.

        :param checks: list of dq rules to save
        :param filepath: path to a file containing the checks.
        """
