import logging
from dataclasses import dataclass
from typing import Any, Callable

import numpy as np

from mlflow.metrics.base import MetricValue
from mlflow.models.evaluation.base import EvaluationMetric

_logger = logging.getLogger(__name__)


@dataclass
class MetricDefinition:
    """
    A dataclass representing a metric definition used in model evaluation.

    Attributes:
        function: The metric function to be called for evaluation.
        name: The name of the metric.
        index: The index of the metric in the ``extra_metrics`` argument of ``mlflow.evaluate``.
        version: (Optional) The metric version. For example v1.
        genai_metric_args: (Optional) A dictionary containing arguments specified by users when
            calling make_genai_metric or make_genai_metric_from_prompt.
            Those args are persisted so that we can deserialize the same metric object later.
    """

    function: Callable[..., Any]
    name: str
    index: int
    version: str | None = None
    genai_metric_args: dict[str, Any] | None = None

    @classmethod
    def from_index_and_metric(cls, index: int, metric: EvaluationMetric):
        return cls(
            function=metric.eval_fn,
            index=index,
            name=metric.name,
            version=metric.version,
            genai_metric_args=metric.genai_metric_args,
        )

    def evaluate(self, eval_fn_args) -> MetricValue | None:
        """
        This function calls the metric function and performs validations on the returned
        result to ensure that they are in the expected format. It will warn and will not log metrics
        that are in the wrong format.

        Args:
            eval_fn_args: A dictionary of args needed to compute the eval metrics.

        Returns:
            MetricValue
        """
        if self.index < 0:
            exception_header = f"Did not log builtin metric '{self.name}' because it"
        else:
            exception_header = (
                f"Did not log metric '{self.name}' at index "
                f"{self.index} in the `extra_metrics` parameter because it"
            )

        metric: MetricValue = self.function(*eval_fn_args)

        def _is_numeric(value):
            return isinstance(value, (int, float, np.number))

        def _is_string(value):
            return isinstance(value, str)

        if metric is None:
            _logger.warning(f"{exception_header} returned None.")
            return

        if _is_numeric(metric):
            return MetricValue(aggregate_results={self.name: metric})

        if not isinstance(metric, MetricValue):
            _logger.warning(f"{exception_header} did not return a MetricValue.")
            return

        scores = metric.scores
        justifications = metric.justifications
        aggregates = metric.aggregate_results

        if scores is not None:
            if not isinstance(scores, list):
                _logger.warning(
                    f"{exception_header} must return MetricValue with scores as a list."
                )
                return
            if any(not (_is_numeric(s) or _is_string(s) or s is None) for s in scores):
                _logger.warning(
                    f"{exception_header} must return MetricValue with numeric or string scores."
                )
                return

        if justifications is not None:
            if not isinstance(justifications, list):
                _logger.warning(
                    f"{exception_header} must return MetricValue with justifications as a list."
                )
                return
            if any(not (_is_string(just) or just is None) for just in justifications):
                _logger.warning(
                    f"{exception_header} must return MetricValue with string justifications."
                )
                return

        if aggregates is not None:
            if not isinstance(aggregates, dict):
                _logger.warning(
                    f"{exception_header} must return MetricValue with aggregate_results as a dict."
                )
                return

            if any(
                not (isinstance(k, str) and (_is_numeric(v) or v is None))
                for k, v in aggregates.items()
            ):
                _logger.warning(
                    f"{exception_header} must return MetricValue with aggregate_results with "
                    "str keys and numeric values."
                )
                return

        return metric