from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    Callable,
    Dict,
    List,
    NamedTuple,
    Optional,
    Tuple,
    Type,
    Union,
)

import great_expectations.exceptions as gx_exceptions
from great_expectations.core.id_dict import IDDict

if TYPE_CHECKING:
    from great_expectations.core.metric_function_types import (
        MetricFunctionTypes,
        MetricPartialFunctionTypes,
    )
    from great_expectations.execution_engine import ExecutionEngine
    from great_expectations.expectations.expectation import Expectation
    from great_expectations.expectations.expectation_configuration import (
        ExpectationConfiguration,
    )
    from great_expectations.expectations.metrics.metric_provider import MetricProvider
    from great_expectations.render import (
        AtomicDiagnosticRendererType,
        AtomicPrescriptiveRendererType,
        AtomicRendererType,
        RenderedAtomicContent,
        RenderedContent,
    )
    from great_expectations.validator.computed_metric import MetricValue

logger = logging.getLogger(__name__)

_registered_expectations: dict = {}
_registered_metrics: dict = {}
_registered_renderers: dict = {}

"""
{
  "metric_name"
    metric_domain_keys
    metric_value_keys
    metric_dependencies
    providers:
      engine: provider
}
"""


class RendererImpl(NamedTuple):
    expectation: str
    renderer: Callable[
        ..., Union[RenderedAtomicContent, list[RenderedAtomicContent], RenderedContent]
    ]


def register_renderer(
    object_name: str,
    parent_class: Union[Type[Expectation], Type[MetricProvider]],
    renderer_fn: Callable[
        ..., Union[RenderedAtomicContent, list[RenderedAtomicContent], RenderedContent]
    ],
):
    # noinspection PyUnresolvedReferences
    renderer_name = renderer_fn._renderer_type  # type: ignore[attr-defined] # FIXME CoP
    if object_name not in _registered_renderers:
        logger.debug(f"Registering {renderer_name} for expectation_type {object_name}.")
        _registered_renderers[object_name] = {renderer_name: (parent_class, renderer_fn)}
        return

    if renderer_name in _registered_renderers[object_name]:
        if _registered_renderers[object_name][renderer_name] == (
            parent_class,
            renderer_fn,
        ):
            logger.info(
                f"Multiple declarations of {renderer_name} renderer for expectation_type {object_name} "  # noqa: E501 # FIXME CoP
                f"found."
            )
            return
        else:
            logger.warning(
                f"Overwriting declaration of {renderer_name} renderer for expectation_type "
                f"{object_name}."
            )
            _registered_renderers[object_name][renderer_name] = (
                parent_class,
                renderer_fn,
            )
        return
    else:
        logger.debug(f"Registering {renderer_name} for expectation_type {object_name}.")
        _registered_renderers[object_name][renderer_name] = (parent_class, renderer_fn)
        return


def get_renderer_names(expectation_or_metric_type: str) -> List[str]:
    """Gets renderer names for a given Expectation or Metric.

    Args:
        expectation_or_metric_type: The type of an Expectation or Metric for which to get renderer names.

    Returns:
        A list of renderer names for the Expectation or Metric.
    """  # noqa: E501 # FIXME CoP
    return list(_registered_renderers.get(expectation_or_metric_type, {}).keys())


def get_renderer_names_with_renderer_types(
    expectation_or_metric_type: str,
    renderer_types: List[AtomicRendererType],
) -> List[Union[str, AtomicDiagnosticRendererType, AtomicPrescriptiveRendererType]]:
    """Gets renderer names of a given type, for a given Expectation or Metric.

    Args:
        expectation_or_metric_type: The type of an Expectation or Metric for which to get renderer names.
        renderer_types: The type of the renderers for which to return names.

    Returns:
        A list of renderer names for the given prefixes and Expectation or Metric.
    """  # noqa: E501 # FIXME CoP
    return [
        renderer_name
        for renderer_name in get_renderer_names(
            expectation_or_metric_type=expectation_or_metric_type
        )
        if any(renderer_name.startswith(renderer_type) for renderer_type in renderer_types)
    ]


def get_renderer_impls(object_name: str) -> List[str]:
    return list(_registered_renderers.get(object_name, {}).values())


def get_renderer_impl(object_name: str, renderer_type: str) -> Optional[RendererImpl]:
    renderer_tuple: Optional[tuple] = _registered_renderers.get(object_name, {}).get(renderer_type)
    renderer_impl: Optional[RendererImpl] = None
    if renderer_tuple:
        renderer_impl = RendererImpl(expectation=renderer_tuple[0], renderer=renderer_tuple[1])
    return renderer_impl


def register_expectation(expectation: Type[Expectation]) -> None:
    expectation_type = expectation.expectation_type
    # TODO: add version to key
    if expectation_type in _registered_expectations:
        if _registered_expectations[expectation_type] == expectation:
            logger.info(f"Multiple declarations of expectation {expectation_type} found.")
            return
        else:
            logger.warning(f"Overwriting declaration of expectation {expectation_type}.")

    logger.debug(f"Registering expectation: {expectation_type}")
    _registered_expectations[expectation_type] = expectation


def register_core_metrics() -> None:
    """As Metric registration is the responsibility of MetaMetricProvider.__new__,
    simply importing a given class will ensure that it is added to the Metric
    registry.

    We use this to grab metrics by name within our workflows.

    Without this function, we need to hope that core Metrics are imported somewhere
    in our import graph - if not, our registry will be empty and we'll see
    MetricResolutionErrors.
    """
    before_count = len(_registered_metrics)

    # Implicitly calls MetaMetricProvider.__new__ as Metrics are loaded from metrics.__init__.py
    # As __new__ calls upon register_metric this import builds our core registry
    from great_expectations.expectations import metrics  # noqa: F401 # FIXME CoP

    after_count = len(_registered_metrics)

    if before_count == after_count:
        logger.debug("Already registered core metrics; no updates to registry")
    else:
        logger.debug(f"Registered {after_count - before_count} core metrics")


def register_core_expectations() -> None:
    """As Expectation registration is the responsibility of MetaExpectation.__new__,
    simply importing a given class will ensure that it is added to the Expectation
    registry.

    We use this JIT in the Validator to ensure that core Expectations are available
    for usage when called upon.

    Without this function, we need to hope that core Expectations are imported somewhere
    in our import graph - if not, our registry will be empty and Validator workflows
    will fail.
    """
    before_count = len(_registered_expectations)

    # Implicitly calls MetaExpectation.__new__ as Expectations are loaded from core.__init__.py
    # As __new__ calls upon register_expectation, this import builds our core registry
    from great_expectations.expectations import core  # noqa: F401 # FIXME CoP

    after_count = len(_registered_expectations)

    if before_count == after_count:
        logger.debug("Already registered core expectations; no updates to registry")
    else:
        logger.debug(f"Registered {after_count - before_count} core expectations")


def _add_response_key(res, key, value):
    if key in res:
        res[key].append(value)
    else:
        res[key] = [value]
    return res


def register_metric(  # noqa: PLR0913 # FIXME CoP
    metric_name: str,
    metric_domain_keys: Tuple[str, ...],
    metric_value_keys: Tuple[str, ...],
    execution_engine: Type[ExecutionEngine],
    metric_class: Type[MetricProvider],
    metric_provider: Optional[Callable],
    metric_fn_type: Optional[Union[MetricFunctionTypes, MetricPartialFunctionTypes]] = None,
) -> dict:
    """Register a Metric class for use as a callable metric within Expectations.

    Args:
        metric_name: A name identifying the metric. Metric Name must be globally unique in
            a great_expectations installation.
        metric_domain_keys: A tuple of the keys used to determine the domain of the metric.
        metric_value_keys: A tuple of the keys used to determine the value of the metric.
        execution_engine: The execution_engine used to execute the metric.
        metric_class: A valid Metric class containing logic to compute attributes of data.
        metric_provider: The MetricProvider class from which the metric_class inherits.
        metric_fn_type: The MetricFunctionType or MetricPartialFunctionType used to define the Metric class.

    Returns:
        A dictionary containing warnings thrown during registration if applicable, and the success status of registration.
    """  # noqa: E501 # FIXME CoP
    res: dict = {}
    execution_engine_name = execution_engine.__name__
    logger.debug(f"Registering metric: {metric_name}")
    if metric_provider is not None and metric_fn_type is not None:
        metric_provider.metric_fn_type = metric_fn_type  # type: ignore[attr-defined] # FIXME CoP
    if metric_name in _registered_metrics:
        metric_definition = _registered_metrics[metric_name]
        current_domain_keys = metric_definition.get("metric_domain_keys", set())
        if set(current_domain_keys) != set(metric_domain_keys):
            logger.warning(
                f"metric {metric_name} is being registered with different metric_domain_keys; overwriting metric_domain_keys"  # noqa: E501 # FIXME CoP
            )
            _add_response_key(
                res,
                "warning",
                f"metric {metric_name} is being registered with different metric_domain_keys; overwriting metric_domain_keys",  # noqa: E501 # FIXME CoP
            )

        current_value_keys = metric_definition.get("metric_value_keys", set())
        if set(current_value_keys) != set(metric_value_keys):
            logger.warning(
                f"metric {metric_name} is being registered with different metric_value_keys; overwriting metric_value_keys"  # noqa: E501 # FIXME CoP
            )
            _add_response_key(
                res,
                "warning",
                f"metric {metric_name} is being registered with different metric_value_keys; overwriting metric_value_keys",  # noqa: E501 # FIXME CoP
            )

        providers = metric_definition.get("providers", {})
        if execution_engine_name in providers:
            _current_provider_cls, current_provider_fn = providers[execution_engine_name]
            if current_provider_fn != metric_provider:
                logger.warning(
                    f"metric {metric_name} is being registered with different metric_provider; overwriting metric_provider"  # noqa: E501 # FIXME CoP
                )
                _add_response_key(
                    res,
                    "warning",
                    f"metric {metric_name} is being registered with different metric_provider; overwriting metric_provider",  # noqa: E501 # FIXME CoP
                )
                providers[execution_engine_name] = metric_class, metric_provider
            else:
                logger.info(
                    f"Multiple declarations of metric {metric_name} for engine {execution_engine_name}."  # noqa: E501 # FIXME CoP
                )
                _add_response_key(
                    res,
                    "info",
                    f"Multiple declarations of metric {metric_name} for engine {execution_engine_name}.",  # noqa: E501 # FIXME CoP
                )
        else:
            providers[execution_engine_name] = metric_class, metric_provider
    else:
        metric_definition = {
            "metric_domain_keys": metric_domain_keys,
            "metric_value_keys": metric_value_keys,
            "default_kwarg_values": metric_class.default_kwarg_values,
            "providers": {execution_engine_name: (metric_class, metric_provider)},
        }
        _registered_metrics[metric_name] = metric_definition

    res["success"] = True

    return res


def _get_metric_definition(metric_name: str) -> dict:
    try:
        return _registered_metrics[metric_name]
    except KeyError:
        raise gx_exceptions.MetricProviderError(f"No metric named {metric_name} found.")  # noqa: TRY003 # FIXME CoP


def get_sqlalchemy_metric_provider(
    metric_name: str,
) -> Tuple[MetricProvider, Callable]:
    """The default SqlAlchemy metric provider for a given metric."""
    metric_definition = _get_metric_definition(metric_name)
    try:
        return metric_definition["providers"]["SqlAlchemyExecutionEngine"]
    except KeyError:
        raise gx_exceptions.MetricProviderError(  # noqa: TRY003 # FIXME CoP
            f"No provider found for {metric_name} using SqlAlchemyExecutionEngine"
        )


def get_metric_provider(
    metric_name: str, execution_engine: ExecutionEngine
) -> Tuple[MetricProvider, Callable]:
    metric_definition = _get_metric_definition(metric_name)
    try:
        return metric_definition["providers"][type(execution_engine).__name__]
    except KeyError:
        # Search up class hierarchy for a match. We skip the first entry since that's the
        # execution engine type itself, type(execution_engine), which we just checked and
        # resulted in the KeyError we're handling here.
        for cls in type(execution_engine).mro()[1:]:
            possible_key = cls.__name__
            if metric_definition["providers"].get(possible_key) is not None:
                metric_provider = metric_definition["providers"][possible_key]
                # Register the metric provider for this engine so we don't have to search again
                metric_definition["providers"][type(execution_engine).__name__] = metric_provider
                return metric_provider
        # no matches when search hierarchy so we raise
        raise gx_exceptions.MetricProviderError(  # noqa: TRY003 # FIXME CoP
            f"No provider found for {metric_name} using {type(execution_engine).__name__}"
        )


def get_metric_function_type(
    metric_name: str, execution_engine: ExecutionEngine
) -> Optional[Union[MetricPartialFunctionTypes, MetricFunctionTypes]]:
    try:
        provider_fn, _provider_class = get_metric_provider(metric_name, execution_engine)
        return getattr(provider_fn, "metric_fn_type", None)
    except KeyError:
        raise gx_exceptions.MetricProviderError(  # noqa: TRY003 # FIXME CoP
            f"No provider found for {metric_name} using {type(execution_engine).__name__}"
        )


def get_metric_kwargs(
    metric_name: str,
    configuration: Optional[ExpectationConfiguration] = None,
    runtime_configuration: Optional[dict] = None,
) -> dict:
    try:
        metric_definition = _registered_metrics.get(metric_name)
        if metric_definition is None:
            raise gx_exceptions.MetricProviderError(f"No definition found for {metric_name}")  # noqa: TRY003 # FIXME CoP
        default_kwarg_values = metric_definition["default_kwarg_values"]
        metric_kwargs = {
            "metric_domain_keys": metric_definition["metric_domain_keys"],
            "metric_value_keys": metric_definition["metric_value_keys"],
        }
        if configuration:
            expectation = configuration.to_domain_obj()
            configuration_kwargs = expectation._get_runtime_kwargs(
                runtime_configuration=runtime_configuration
            )
            if len(metric_kwargs["metric_domain_keys"]) > 0:
                metric_domain_kwargs = IDDict(
                    {
                        k: configuration_kwargs.get(k) or default_kwarg_values.get(k)
                        for k in metric_kwargs["metric_domain_keys"]
                    }
                )
            else:
                metric_domain_kwargs = IDDict()
            if len(metric_kwargs["metric_value_keys"]) > 0:
                metric_value_kwargs = IDDict(
                    {
                        k: configuration_kwargs.get(k)
                        if configuration_kwargs.get(k) is not None
                        else default_kwarg_values.get(k)
                        for k in metric_kwargs["metric_value_keys"]
                    }
                )
            else:
                metric_value_kwargs = IDDict()
            metric_kwargs["metric_domain_kwargs"] = metric_domain_kwargs
            metric_kwargs["metric_value_kwargs"] = metric_value_kwargs
        return metric_kwargs
    except KeyError:
        raise gx_exceptions.MetricProviderError(f"Incomplete definition found for {metric_name}")  # noqa: TRY003 # FIXME CoP


def get_domain_metrics_dict_by_name(
    metrics: Dict[Tuple[str, str, str], MetricValue], metric_domain_kwargs: IDDict
):
    return {
        metric_edge_key_id_tuple[0]: metric_value
        for metric_edge_key_id_tuple, metric_value in metrics.items()
        if metric_edge_key_id_tuple[1] == metric_domain_kwargs.to_id()
    }


def get_expectation_impl(expectation_name: str) -> Type[Expectation]:
    expectation: Type[Expectation] | None = _registered_expectations.get(expectation_name)
    if not expectation:
        raise gx_exceptions.ExpectationNotFoundError(f"{expectation_name} not found")  # noqa: TRY003 # FIXME CoP

    return expectation


def list_registered_expectation_implementations(
    expectation_root: Optional[Type[Expectation]] = None,
) -> List[str]:
    registered_expectation_implementations = []
    for (
        expectation_name,
        expectation_implementation,
    ) in _registered_expectations.items():
        if expectation_root is None or (
            expectation_root and issubclass(expectation_implementation, expectation_root)
        ):
            registered_expectation_implementations.append(expectation_name)

    return registered_expectation_implementations
