import logging
import json
import warnings
from typing import Any
from collections.abc import Callable
from pathlib import Path

import yaml
from pyspark.sql import DataFrame, SparkSession

from databricks.labs.dqx.checks_resolver import resolve_check_function
from databricks.labs.dqx.checks_validator import ChecksValidator
from databricks.labs.dqx.rule import (
    DQRule,
    DQRowRule,
    DQDatasetRule,
    DQForEachColRule,
    CHECK_FUNC_REGISTRY,
)
from databricks.labs.dqx.utils import safe_json_load
from databricks.labs.dqx.errors import InvalidCheckError

CHECKS_TABLE_SCHEMA = (
    "name STRING, criticality STRING, check STRUCT<function STRING, for_each_column ARRAY<STRING>,"
    " arguments MAP<STRING, STRING>>, filter STRING, run_config_name STRING, user_metadata MAP<STRING, STRING>"
)

FILE_SERIALIZERS: dict[str, Callable[[list[dict[Any, Any]]], str]] = {
    ".json": json.dumps,
    ".yml": yaml.safe_dump,
    ".yaml": yaml.safe_dump,
}


FILE_DESERIALIZERS: dict[str, Callable] = {
    ".json": json.load,
    ".yaml": yaml.safe_load,
    ".yml": yaml.safe_load,
}

logger = logging.getLogger(__name__)


def serialize_checks_from_dataframe(df: DataFrame, run_config_name: str = "default") -> list[dict]:
    """
    Converts a list of quality checks defined in a DataFrame to a list of quality checks
    defined as Python dictionaries.

    Args:
        df: DataFrame with data quality check rules. Each row should define a check. Rows should
        have the following columns:
            - *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).
            - *check* - DQX check function used in the check; A *StructType* column defining the data quality check.
            - *filter* - Expression for filtering data quality checks.
            - *run_config_name* (optional) - Run configuration name for storing checks across runs.
            - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check.
        run_config_name: Run configuration name for filtering quality rules, e.g. input table or job name (use "default" if not provided).

    Returns:
            List of data quality check specifications as a Python dictionary
    """
    check_rows = df.where(f"run_config_name = '{run_config_name}'").collect()
    collect_limit = 500
    if len(check_rows) > collect_limit:
        warnings.warn(
            f"Collecting large number of rows from DataFrame: {len(check_rows)}",
            category=UserWarning,
            stacklevel=2,
        )

    checks = []
    for row in check_rows:
        check_dict = {
            "name": row.name,
            "criticality": row.criticality,
            "check": {
                "function": row.check["function"],
                "arguments": (
                    {k: safe_json_load(v) for k, v in row.check["arguments"].items()}
                    if row.check["arguments"] is not None
                    else {}
                ),
            },
        }
        if "for_each_column" in row.check and row.check["for_each_column"]:
            check_dict["check"]["for_each_column"] = row.check["for_each_column"]
        if row.filter is not None:
            check_dict["filter"] = row.filter
        if row.user_metadata is not None:
            check_dict["user_metadata"] = row.user_metadata
        checks.append(check_dict)
    return checks


def deserialize_checks_to_dataframe(
    spark: SparkSession,
    checks: list[dict],
    run_config_name: str = "default",
) -> DataFrame:
    """
    Converts a list of quality checks defined as Python dictionaries to a DataFrame.

    Args:
        spark: Spark session.
        checks: list of check specifications as Python dictionaries. Each check consists of the 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)
            - *filter* (optional) - Expression for filtering data quality checks
            - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check.
        run_config_name: Run configuration name for storing quality checks across runs, e.g. input table or job name (use "default" if not provided)

    Returns:
        DataFrame with data quality check rules

    Raises:
        InvalidCheckError: If any check is invalid or unsupported.
    """
    dq_rule_checks: list[DQRule] = deserialize_checks(checks)

    dq_rule_rows = []
    for dq_rule_check in dq_rule_checks:
        arguments = dq_rule_check.check_func_kwargs

        if dq_rule_check.column is not None:
            arguments["column"] = dq_rule_check.column

        if dq_rule_check.columns is not None:
            arguments["columns"] = dq_rule_check.columns

        # row_filter is resolved from the check filter so not need to include
        json_arguments = {k: json.dumps(v) for k, v in arguments.items() if k not in {"row_filter"}}
        dq_rule_rows.append(
            [
                dq_rule_check.name,
                dq_rule_check.criticality,
                {"function": dq_rule_check.check_func.__name__, "arguments": json_arguments},
                dq_rule_check.filter,
                run_config_name,
                dq_rule_check.user_metadata,
            ]
        )
    return spark.createDataFrame(dq_rule_rows, CHECKS_TABLE_SCHEMA)


def deserialize_checks(checks: list[dict], custom_checks: dict[str, Callable] | None = None) -> list[DQRule]:
    """
    Converts a list of quality checks defined as Python dictionaries to a list of `DQRule` objects.

    Args:
        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
                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)
            - *filter* (optional) - Expression for filtering data quality checks
            - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check.
        custom_checks: dictionary with custom check functions (e.g. *globals()* of the calling module).
            If not specified, then only built-in functions are used for the checks.

    Returns:
        list of data quality check rules

    Raises:
        InvalidCheckError: If any dictionary is invalid or unsupported.
    """
    status = ChecksValidator.validate_checks(checks, custom_checks)
    if status.has_errors:
        raise InvalidCheckError(str(status))

    dq_rule_checks: list[DQRule] = []
    for check_def in checks:
        logger.debug(f"Processing check definition: {check_def}")

        check = check_def.get("check", {})
        name = check_def.get("name", None)
        func_name = check.get("function")
        func = resolve_check_function(func_name, custom_checks, fail_on_missing=True)
        assert func  # should already be validated

        func_args = check.get("arguments", {})
        for_each_column = check.get("for_each_column")
        column = func_args.get("column")  # should be defined for single-column checks only
        columns = func_args.get("columns")  # should be defined for multi-column checks only
        assert not (column and columns)  # should already be validated
        criticality = check_def.get("criticality", "error")
        filter_str = check_def.get("filter")
        user_metadata = check_def.get("user_metadata")

        # Exclude `column` and `columns` from check_func_kwargs
        # as these are always included in the check function call
        check_func_kwargs = {k: v for k, v in func_args.items() if k not in {"column", "columns"}}

        # treat non-registered function as row-level checks
        if for_each_column:
            dq_rule_checks += DQForEachColRule(
                columns=for_each_column,
                name=name,
                check_func=func,
                criticality=criticality,
                filter=filter_str,
                check_func_kwargs=check_func_kwargs,
                user_metadata=user_metadata,
            ).get_rules()
        else:
            rule_type = CHECK_FUNC_REGISTRY.get(func_name)
            if rule_type == "dataset":
                dq_rule_checks.append(
                    DQDatasetRule(
                        column=column,
                        columns=columns,
                        check_func=func,
                        check_func_kwargs=check_func_kwargs,
                        name=name,
                        criticality=criticality,
                        filter=filter_str,
                        user_metadata=user_metadata,
                    )
                )
            else:  # default to row-level rule
                dq_rule_checks.append(
                    DQRowRule(
                        column=column,
                        columns=columns,
                        check_func=func,
                        check_func_kwargs=check_func_kwargs,
                        name=name,
                        criticality=criticality,
                        filter=filter_str,
                        user_metadata=user_metadata,
                    )
                )

    return dq_rule_checks


def serialize_checks(checks: list[DQRule]) -> list[dict]:
    """
    Converts a list of quality checks defined as *DQRule* objects to a list of quality checks
    defined as Python dictionaries.

    Args:
        checks: List of DQRule instances to convert.

    Returns:
        List of dictionaries representing the DQRule instances.

    Raises:
        InvalidCheckError: If any item in the list is not a DQRule instance.
    """
    dq_rules = []
    for check in checks:
        if not isinstance(check, DQRule):
            raise InvalidCheckError(f"Expected DQRule instance, got {type(check).__name__}")
        dq_rules.append(check.to_dict())
    return dq_rules


def serialize_checks_to_bytes(checks: list[dict], file_path: Path) -> bytes:
    """
    Serializes a list of checks to bytes in json or yaml (default) format.

    Args:
        checks: List of checks to serialize.
        file_path: Path to the file where the checks will be serialized.
    Returns:
        Serialized checks as bytes.
    """
    serializer = FILE_SERIALIZERS.get(file_path.suffix.lower(), yaml.safe_dump)  # default to yaml
    return serializer(checks).encode("utf-8")


def get_file_deserializer(filepath: str) -> Callable:
    """
    Get the deserializer function based on file.

    Args:
        filepath: Path to the file.
    Returns:
        Deserializer function.
    """
    ext = Path(filepath).suffix.lower()
    return FILE_DESERIALIZERS.get(ext.lower(), yaml.safe_load)  # default to yaml
