from __future__ import annotations

import logging
import traceback
from typing import (
    TYPE_CHECKING,
    Callable,
    Dict,
    List,
    Optional,
    Set,
    Tuple,
    TypedDict,
    Union,
)

from tqdm.auto import tqdm

import great_expectations.exceptions as gx_exceptions
from great_expectations.compatibility.typing_extensions import override
from great_expectations.expectations.registry import get_metric_provider
from great_expectations.validator.exception_info import ExceptionInfo
from great_expectations.validator.metric_configuration import (
    MetricConfiguration,
    MetricConfigurationID,
)

if TYPE_CHECKING:
    from great_expectations.core import IDDict
    from great_expectations.execution_engine import ExecutionEngine
    from great_expectations.expectations.expectation_configuration import (
        ExpectationConfiguration,
    )
    from great_expectations.expectations.metrics.metric_provider import MetricProvider
    from great_expectations.validator.computed_metric import MetricValue
    from great_expectations.validator.metrics_calculator import (
        _AbortedMetricsInfoDict,
    )


class MetricsCalculatorErrorResultValue(TypedDict):
    metric_configuration: MetricConfiguration
    exception_info: ExceptionInfo
    num_failures: int


__all__ = [
    "ExpectationValidationGraph",
    "MetricConfiguration",
    "MetricEdge",
    "MetricValue",
    "ValidationGraph",
]

logger = logging.getLogger(__name__)
logging.captureWarnings(True)

MAX_METRIC_COMPUTATION_RETRIES: int = 3


class MetricEdge:
    def __init__(
        self, left: MetricConfiguration, right: Optional[MetricConfiguration] = None
    ) -> None:
        self._left = left
        self._right = right

    @property
    def left(self):
        return self._left

    @property
    def right(self):
        return self._right

    @property
    def id(self):
        if self.right:
            return self.left.id, self.right.id

        return self.left.id, None

    def __repr__(self):  # type: ignore[explicit-override] # FIXME
        return f"<{self._left.__repr__()}|{self._right.__repr__()}>"


class ValidationGraph:
    def __init__(
        self,
        execution_engine: ExecutionEngine,
        edges: Optional[List[MetricEdge]] = None,
    ) -> None:
        self._execution_engine = execution_engine

        if edges:
            self._edges = edges
        else:
            self._edges = []

        self._edge_ids = {edge.id for edge in self._edges}

    @override
    def __eq__(self, other) -> bool:
        """Supports comparing two "ValidationGraph" objects."""
        return self.edge_ids == other.edge_ids

    @override
    def __hash__(self) -> int:
        return hash(frozenset(self.edge_ids))

    @property
    def edges(self) -> List[MetricEdge]:
        """Returns "MetricEdge" objects, contained within this "ValidationGraph" object (as list)."""  # noqa: E501 # FIXME CoP
        return self._edges

    @property
    def edge_ids(self) -> Set[Tuple[str, str]]:
        """Returns "MetricEdge" objects, contained within this "ValidationGraph" object (as set of two-tuples)."""  # noqa: E501 # FIXME CoP
        return {edge.id for edge in self._edges}

    def add(self, edge: MetricEdge) -> None:
        """Adds supplied "MetricEdge" object to this "ValidationGraph" object (if not already present)."""  # noqa: E501 # FIXME CoP
        if edge.id not in self._edge_ids:
            self._edges.append(edge)
            self._edge_ids.add(edge.id)

    def build_metric_dependency_graph(
        self,
        metric_configuration: MetricConfiguration,
        runtime_configuration: Optional[dict] = None,
    ) -> None:
        """
        Obtain domain and value keys for metrics and proceeds to add these metrics to the validation graph
        until all metrics have been added.

        Args:
            metric_configuration: Desired MetricConfiguration object to be resolved.
            runtime_configuration: Additional run-time settings (see "Validator.DEFAULT_RUNTIME_CONFIGURATION").
        """  # noqa: E501 # FIXME CoP

        metric_impl_klass: MetricProvider
        metric_provider: Callable
        (
            metric_impl_klass,
            _metric_provider,
        ) = self.set_metric_configuration_default_kwargs_if_absent(
            metric_configuration=metric_configuration
        )

        metric_dependencies = metric_impl_klass.get_evaluation_dependencies(
            metric=metric_configuration,
            execution_engine=self._execution_engine,
            runtime_configuration=runtime_configuration,
        )

        if len(metric_dependencies) == 0:
            self.add(
                MetricEdge(
                    left=metric_configuration,
                )
            )
        else:
            metric_configuration.metric_dependencies = metric_dependencies
            for metric_dependency in metric_dependencies.values():
                # TODO: <Alex>In the future, provide a more robust cycle detection mechanism.</Alex>
                if metric_dependency.id == metric_configuration.id:
                    logger.warning(
                        f"Metric {metric_configuration.id!s} has created a circular dependency"
                    )
                    continue
                self.add(
                    MetricEdge(
                        left=metric_configuration,
                        right=metric_dependency,
                    )
                )
                self.build_metric_dependency_graph(
                    metric_configuration=metric_dependency,
                    runtime_configuration=runtime_configuration,
                )

    def set_metric_configuration_default_kwargs_if_absent(
        self, metric_configuration: MetricConfiguration
    ) -> Tuple[MetricProvider, Callable]:
        """
        Updates "metric_domain_kwargs" and/or "metric_value_kwargs" of "MetricConfiguration" with defualts (if needed).
        """  # noqa: E501 # FIXME CoP
        metric_impl_klass: MetricProvider
        metric_provider: Callable
        metric_impl_klass, metric_provider = get_metric_provider(
            metric_name=metric_configuration.metric_name,
            execution_engine=self._execution_engine,
        )
        self._set_default_metric_kwargs_if_absent(
            default_kwarg_values=metric_impl_klass.default_kwarg_values,
            metric_kwargs=metric_configuration.metric_domain_kwargs,
            keys=metric_impl_klass.domain_keys,
        )
        self._set_default_metric_kwargs_if_absent(
            default_kwarg_values=metric_impl_klass.default_kwarg_values,
            metric_kwargs=metric_configuration.metric_value_kwargs,
            keys=metric_impl_klass.value_keys,
        )
        return metric_impl_klass, metric_provider

    def resolve(
        self,
        runtime_configuration: Optional[dict] = None,
        min_graph_edges_pbar_enable: int = 0,
        # Set to low number (e.g., 3) to suppress progress bar for small graphs.
        show_progress_bars: bool = True,
    ) -> Tuple[
        Dict[MetricConfigurationID, MetricValue],
        _AbortedMetricsInfoDict,
    ]:
        resolved_metrics: Dict[MetricConfigurationID, MetricValue] = {}

        # updates graph with aborted metrics
        aborted_metrics_info: _AbortedMetricsInfoDict = self._resolve(
            metrics=resolved_metrics,
            runtime_configuration=runtime_configuration,
            min_graph_edges_pbar_enable=min_graph_edges_pbar_enable,
            show_progress_bars=show_progress_bars,
        )

        return resolved_metrics, aborted_metrics_info

    def _resolve(  # noqa: C901, PLR0912, PLR0915 # FIXME CoP
        self,
        metrics: Dict[MetricConfigurationID, MetricValue],
        runtime_configuration: Optional[dict] = None,
        min_graph_edges_pbar_enable: int = 0,  # Set to low number (e.g., 3) to suppress progress bar for small graphs.  # noqa: E501 # FIXME CoP
        show_progress_bars: bool = True,
    ) -> _AbortedMetricsInfoDict:
        if metrics is None:
            metrics = {}

        if runtime_configuration is None:
            runtime_configuration = {}

        if runtime_configuration.get("catch_exceptions", True):
            catch_exceptions = True
        else:
            catch_exceptions = False

        failed_metric_info: _AbortedMetricsInfoDict = {}
        aborted_metrics_info: _AbortedMetricsInfoDict = {}

        ready_metrics: Set[MetricConfiguration]
        needed_metrics: Set[MetricConfiguration]

        exception_info: ExceptionInfo

        progress_bar: Optional[tqdm] = None

        done: bool = False
        while not done:
            ready_metrics, needed_metrics = self._parse(metrics=metrics)

            # Check to see if the user has disabled progress bars
            disable = not show_progress_bars
            if len(self.edges) < min_graph_edges_pbar_enable:
                disable = True

            if progress_bar is None:
                # noinspection PyProtectedMember,SpellCheckingInspection
                progress_bar = tqdm(
                    total=len(ready_metrics) + len(needed_metrics),
                    desc="Calculating Metrics",
                    disable=disable,
                )
            progress_bar.update(0)
            progress_bar.refresh()

            computable_metrics = set()

            for metric in ready_metrics:
                if (
                    metric.id in failed_metric_info
                    and failed_metric_info[metric.id]["num_failures"]
                    >= MAX_METRIC_COMPUTATION_RETRIES
                ):
                    aborted_metrics_info[metric.id] = failed_metric_info[metric.id]
                else:
                    computable_metrics.add(metric)

            try:
                # Access "ExecutionEngine.resolve_metrics()" method, to resolve missing "MetricConfiguration" objects.  # noqa: E501 # FIXME CoP
                metrics.update(
                    self._execution_engine.resolve_metrics(
                        metrics_to_resolve=computable_metrics,
                        metrics=metrics,
                        runtime_configuration=runtime_configuration,
                    )
                )
                progress_bar.update(len(computable_metrics))
                progress_bar.refresh()
            except gx_exceptions.MetricResolutionError as err:
                if catch_exceptions:
                    exception_traceback = traceback.format_exc()
                    exception_message = str(err)
                    exception_info = ExceptionInfo(
                        exception_traceback=exception_traceback,
                        exception_message=exception_message,
                    )
                    for failed_metric in err.failed_metrics:
                        if failed_metric.id in failed_metric_info:
                            failed_metric_info[failed_metric.id]["num_failures"] += 1
                            failed_metric_info[failed_metric.id]["exception_info"] = exception_info
                        else:
                            failed_metric_info[failed_metric.id] = (
                                MetricsCalculatorErrorResultValue(
                                    metric_configuration=failed_metric,
                                    exception_info=exception_info,
                                    num_failures=1,
                                )
                            )

                else:
                    raise err  # noqa: TRY201 # FIXME CoP
            except Exception as e:
                if catch_exceptions:
                    logger.error(  # noqa: TRY400 # FIXME CoP
                        f"""Caught exception {e!s} while trying to resolve a set of {len(ready_metrics)} metrics; aborting graph resolution."""  # noqa: E501 # FIXME CoP
                    )
                    done = True
                else:
                    raise e  # noqa: TRY201 # FIXME CoP

            if (len(ready_metrics) + len(needed_metrics) == 0) or (
                len(ready_metrics) == len(aborted_metrics_info)
            ):
                done = True

        progress_bar.close()  # type: ignore[union-attr]  # Incorrect flagging of 'Item "None" of "Optional[Any]" has no attribute "close"' in external package.

        return aborted_metrics_info

    def _parse(
        self,
        metrics: Dict[MetricConfigurationID, MetricValue],
    ) -> Tuple[Set[MetricConfiguration], Set[MetricConfiguration]]:
        """Given validation graph, returns the ready and needed metrics necessary for validation using a traversal of
        validation graph (a graph structure of metric ids) edges"""  # noqa: E501 # FIXME CoP
        unmet_dependency_ids = set()
        unmet_dependency = set()
        maybe_ready_ids = set()
        maybe_ready = set()

        for edge in self.edges:
            if edge.left.id not in metrics:
                if edge.right is None or edge.right.id in metrics:
                    if edge.left.id not in maybe_ready_ids:
                        maybe_ready_ids.add(edge.left.id)
                        maybe_ready.add(edge.left)
                else:  # noqa: PLR5501 # FIXME CoP
                    if edge.left.id not in unmet_dependency_ids:
                        unmet_dependency_ids.add(edge.left.id)
                        unmet_dependency.add(edge.left)

        return maybe_ready - unmet_dependency, unmet_dependency

    @staticmethod
    def _set_default_metric_kwargs_if_absent(
        default_kwarg_values: dict,
        metric_kwargs: IDDict,
        keys: Tuple[str, ...],
    ) -> None:
        key: str
        for key in keys:
            if (
                key not in metric_kwargs
                and key in default_kwarg_values
                and default_kwarg_values[key] is not None
            ):
                metric_kwargs[key] = default_kwarg_values[key]

    def __repr__(self):  # type: ignore[explicit-override] # FIXME
        edge: MetricEdge
        return ", ".join([edge.__repr__() for edge in self._edges])


class ExpectationValidationGraph:
    def __init__(
        self,
        configuration: ExpectationConfiguration,
        graph: ValidationGraph,
    ) -> None:
        if configuration is None:
            raise ValueError(  # noqa: TRY003 # FIXME CoP
                """Instantiation of "ExpectationValidationGraph" requires valid "ExpectationConfiguration" object."""  # noqa: E501 # FIXME CoP
            )

        if graph is None:
            raise ValueError(  # noqa: TRY003 # FIXME CoP
                """Instantiation of "ExpectationValidationGraph" requires valid "ValidationGraph" object."""  # noqa: E501 # FIXME CoP
            )

        self._configuration = configuration
        self._graph = graph

    @property
    def configuration(self) -> ExpectationConfiguration:
        return self._configuration

    @property
    def graph(self) -> ValidationGraph:
        return self._graph

    def update(self, graph: ValidationGraph) -> None:
        edge: MetricEdge
        for edge in graph.edges:
            self.graph.add(edge=edge)

    def get_exception_info(
        self,
        metric_info: _AbortedMetricsInfoDict,
    ) -> Dict[str, Union[MetricConfiguration, ExceptionInfo, int]]:
        metric_info = self._filter_metric_info_in_graph(metric_info=metric_info)
        metric_exception_info: Dict[str, Union[MetricConfiguration, ExceptionInfo, int]] = {}
        metric_id: MetricConfigurationID
        metric_info_item: MetricsCalculatorErrorResultValue
        for metric_id, metric_info_item in metric_info.items():
            metric_exception_info[str(metric_id)] = metric_info_item["exception_info"]

        return metric_exception_info

    def _filter_metric_info_in_graph(
        self,
        metric_info: _AbortedMetricsInfoDict,
    ) -> _AbortedMetricsInfoDict:
        graph_metric_ids: List[MetricConfigurationID] = []
        edge: MetricEdge
        vertex: MetricConfiguration
        for edge in self.graph.edges:
            for vertex in [edge.left, edge.right]:
                if vertex is not None:
                    graph_metric_ids.append(vertex.id)

        metric_id: MetricConfigurationID
        metric_info_item: Dict[str, Union[MetricConfiguration, Set[ExceptionInfo], int]]
        return {
            metric_id: metric_info_item
            for metric_id, metric_info_item in metric_info.items()
            if metric_id in graph_metric_ids
        }

    @override
    def __eq__(self, other: object) -> bool:
        if not isinstance(other, ExpectationValidationGraph):
            return False
        return self.configuration == other.configuration and self.graph == other.graph

    @override
    def __hash__(self) -> int:
        return hash((self.configuration, self.graph))
