from typing import Any, Dict, Set, List, Tuple

from great_expectations.core import ExpectationSuite, ExpectationConfiguration, \
    ExpectationValidationResult
from great_expectations.core.batch import Batch
from great_expectations.execution_engine import PandasExecutionEngine
from great_expectations.expectations.expectation import Expectation
from great_expectations.profile.base import ProfilerDataType
import pandas as pd

from datasmelldetection.core.datasmells import DataSmellType
from datasmelldetection.detectors.great_expectations.datasmell import (
    DataSmellRegistry,
    DataSmellMetadata,
)
from great_expectations.validator.validator import Validator

from .helper_dataclasses import DataSmellInformation

# Ensure that a data smell has been registered. Check the following:
# + A dictionary is returned by the
#   get_smell_dict_for_profiler_data_type() call.
# + The data smell type in the metadata object is present in the dictionary
#   returned by get_smell_dict_for_profiler_data_type() as a key.
# + The expectation_type is the corresponding dictionary value.
def check_data_smell_stored_in_registry(
        registry: DataSmellRegistry,
        metadata: DataSmellMetadata,
        expectation_type: str):
    for data_type in metadata.profiler_data_types:
        smell_dict = registry.get_smell_dict_for_profiler_data_type(data_type)
        assert isinstance(smell_dict, dict)

        # Ensure that the data smell has been registered for the corresponding
        # ProfilerDataType.
        assert metadata.data_smell_type in smell_dict
        # Ensure that the expectation type is stored as the corresponding dict
        # value.
        assert smell_dict[metadata.data_smell_type] == expectation_type

    # Data types for which the smell should not be registered.
    not_specified_datatypes = set([
        x for x in ProfilerDataType if x not in metadata.profiler_data_types
    ])
    # Ensure that the data smell was not registered for the not specified data types.
    for data_type in not_specified_datatypes:
        smell_dict = registry.get_smell_dict_for_profiler_data_type(data_type)
        assert metadata.data_smell_type not in smell_dict



# Ensure that profiler data types for which no data smells have been registered
# do not contain any smells. Pass the registry to check and the expected types
# (profiler data types for which data smells have been registered).
def check_remaining_data_types_have_no_registered_smells(
        registry,
        expected_types):
    # ProfilerDataTypes which must not contain registered data smells.
    remaining_types = [
        x for x in ProfilerDataType
        if x not in expected_types
    ]
    for data_type in remaining_types:
        # Ensure that no smells have been registered.
        result = registry.get_smell_dict_for_profiler_data_type(data_type)
        assert isinstance(result, dict)
        assert len(result) == 0


# Ensure that a specific expectation is present in an expectation suite.
# Return True if it was found and False otherwise.
def check_expectation_exists(
        expectation_suite: ExpectationSuite,
        expectation_type: str,
        expected_kwargs: Dict[str, Any]) -> bool:
    for expectation_configuration in expectation_suite.expectations:
        if expectation_configuration.expectation_type == expectation_type and \
                expectation_configuration.kwargs == expected_kwargs:
            # Matching expectation configuration is present.
            return True

    # Matching expectation configuration not found
    return False


# Ensure that all combinations of columns and matching data smells are
# generated. For instance, integer-specific expectations should be generated
# for all integer columns.
#
# expectation_suite: The ExpectationSuite generated by profiling.
# expected_column_types: Information about which profiler data type (column
#   type) has which column names associated. This is specific to a dataset
#   used for profiling.
# data_smell_information: The data smells which the profiler should consider.
def check_all_expectation_combinations_in_expectations_suite(
        expectation_suite: ExpectationSuite,
        expected_column_types: Dict[ProfilerDataType, Set[str]],
        data_smell_information: List[DataSmellInformation]
    ):

    # Ensure that a specific expectation type is present in an expectation
    # suite in certain columns.
    def check_all_columns_present(expectation_type: str, columns: Set[str], kwargs: Dict[str, Any]):
        for column in columns:
            assert check_expectation_exists(
                expectation_suite=expectation_suite,
                expectation_type=expectation_type,
                expected_kwargs={**kwargs, "column": column}
            )

    for information in data_smell_information:
        # Iterate over all column types for which the data smell should be
        # detected.
        for data_type in information.metadata.profiler_data_types:
            if data_type in expected_column_types:
                # The profiled dataset contains columns of such type.
                # => Ensure that an expectation for each corresponding column
                # is generated.
                check_all_columns_present(
                    expectation_type=information.expectation_type,
                    columns=expected_column_types[data_type],
                    kwargs=information.kwargs
                )

# Ensure that the dictionary returned by the
# get_expectation_type_to_data_smell_type_dict method of DataSmellRegistry
# contains the expected mapping. That dictionary stores mappings from the
# expectation type to the corresponding data smell type. This mapping is
# required to reconstruct the data smell type from an expectation type
# as present in ExpectationValidationResult objects which only contain
# the expectation types.
def check_get_expectation_type_to_data_smell_type_dict(
        returned_dict: Dict[str, DataSmellType],
        data_smell_information: List[DataSmellInformation]):
    # expectation_type_to_data_smell_type dictionary which is expected.
    expected_dict = {}

    for information in data_smell_information:
        expectation_type = information.expectation_type
        assert isinstance(expectation_type, str)

        data_smell_type = information.metadata.data_smell_type
        assert isinstance(data_smell_type, DataSmellType)

        expected_dict[expectation_type] = data_smell_type

    assert isinstance(returned_dict, dict)
    assert returned_dict == expected_dict

# Ensure that the get_registered_data_smells method of DataSmellRegistry
# matches the expected set of registered data smells.
def check_get_registered_data_smells(
        returned_set: Set[DataSmellType],
        data_smell_information: List[DataSmellInformation]):
    expected_set: Set[DataSmellType] = set()

    for information in data_smell_information:
        data_smell_type = information.metadata.data_smell_type
        assert isinstance(data_smell_type, DataSmellType)
        expected_set.add(data_smell_type)

    assert isinstance(returned_set, set)
    assert returned_set == expected_set

# To reconstruct the column type information of a profiled dataset later on, it
# is passed in the meta dictionary by the DataSmellAwareProfiler.
def check_column_types_in_expectation_suite_meta_information(
        suite: ExpectationSuite,
        expected_column_types: Dict[ProfilerDataType, Set[str]]):

    # Reconstruct the column type information from the expected_column_types
    # parameter. This variable stores the mapping of column names (string)
    # to a dictionary with information about the corresponding column.
    # The `type` key of the column-specific dictionary stores the corresponding
    # profiler data type in string form. The string representation is used
    # to make the corresponding type JSON-serializable without changing the
    # implementation of the ProfilerDataType enum (in Great Expectations).
    expected_dict: Dict[str, Dict[str, str]] = {}
    for data_type, column_names in expected_column_types.items():
        for column_name in column_names:
            expected_dict[column_name] = {"type": str(data_type)}

    # Ensure that a columns dictionary was generated in the meta information.
    assert "columns" in suite.meta
    result_dict: Dict[str, ProfilerDataType] = suite.meta["columns"]
    assert isinstance(result_dict, dict)
    # Ensure that the generated dictionary matches the expected dictionary.
    assert expected_dict == result_dict


# Execute the examples which the passed expectation contains. Ensure
# a test report exists and that all tests were successful.
#
# NOTE: Internally the examples of a Great Expectations Expectation is checked
# using the mechanism in the run_diagnostics method of the Expectation class.
# The implementation of the 0.13.16 version has been used. This step was
# necessary since the examples don't have to be checked for spark or sqlalchemy.
def check_expectation_examples(expectation: Expectation):
    # Ensure examples are present which should be checked.
    assert hasattr(expectation, "examples")
    examples: List[Dict[str, Any]] = expectation.examples
    assert isinstance(examples, list)
    assert len(examples) == 1, f"Expect one example dataset with tests" \
                               f" for {expectation.expectation_type}"
    expectation_type = expectation.expectation_type

    # Use the only existing test dataset for the corresponding expectation.
    # NOTE: Inlined from _choose_example method of Expectation class.
    example_data = examples[0]["data"]
    for example_test in examples[0]["tests"]:
        example_title: str = example_test["title"]
        print(f"\tExecuting testcase {example_title}")
        # Construct a string which describes for which expectation and
        # which example testcase a test fails.
        example_identifier: str = f"{expectation_type}-{example_title}"

        test_batch = Batch(data=pd.DataFrame(example_data))

        expectation_config = ExpectationConfiguration(
            **{
                "expectation_type": expectation_type,
                "kwargs": example_test["in"]
            }
        )

        # NOTE: Inlined from _instantiate_example_validation_results method of
        # Expectation class.
        validation_results: List[ExpectationValidationResult] = Validator(
            execution_engine=PandasExecutionEngine(),
            batches=[test_batch]
        ).graph_validate(configurations=[expectation_config])

        # NOTE: From run_diagnostics method of Expectation class.
        assert len(validation_results) == 1, example_identifier
        validation_result: ExpectationValidationResult = validation_results[0]

        expected_success = example_test["out"]["success"]
        assert validation_result.success == expected_success, example_identifier

        # Construct a list of keys for which comparison between keys in the
        # validation result object and the example data should be performed.
        # Success is not compared since checking has been performed above.
        out_keys = example_test["out"].keys()
        comparison_keys = [x for x in out_keys if x != "success"]

        # Perform pairwise comparison of validation results and example test
        # data.
        for key in comparison_keys:
            assert key in validation_result.result, \
                f"{example_identifier} no result for key {key}."
            expected = set(example_test["out"][key])
            actual = set(validation_result.result[key])
            # Use set equality since the order for results like
            # "partial_unexpected_list" does not matter.
            assert actual == expected, f"{example_identifier}: Failed for key {key}"


# Only keep columns which are present in the specified column names.
# This helper function is required to test column name filtering of the
# DataSmellAwareProfiler.
def filter_expected_column_types(
        expected_column_types: Dict[ProfilerDataType, Set[str]],
        column_names: Set[str]) -> Dict[ProfilerDataType, Set[str]]:
    result: Dict[ProfilerDataType, Set[str]] = dict()

    # Only keep relevant column entries.
    for data_type, columns in expected_column_types.items():
        result[data_type] = {x for x in columns if x in column_names}

    return result
