#!/usr/bin/env python3

"""Abstract generic base class `BaseMetric[T]` and its `DefaultSchema` configuration

![status](../../../badges/status-metrics-base.svg)
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections import deque
from collections.abc import Iterator
from types import get_original_bases
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Generic,
    Literal,
    TypeVar,
    cast,
    get_args,
    get_origin,
)

import orjson
from pydantic import BaseModel, ConfigDict, ValidationError

from evallm.metrics.types import (
    SerializableValueType,
    UnserializableValueTypeError,
    check_evaluate_signature,
    is_serializable_type,
    matches_annotation,
)

if TYPE_CHECKING:
    from evallm.metrics.composite import CompositeMetric


T = TypeVar("T", bound=SerializableValueType)


class DefaultSchema(BaseModel):
    """Default empty schema as a `pydantic.BaseModel`.

    Used as the default value for the `config_schema` class variable in `BaseMetric`
    to match the shared interface and indicate that a metric doesn't take any
    configuration parameters. Schema enforces that no configuration parameters can be
    passed at instantiation by forbidding extra fields and having no defined fields. (
    see also: [`BaseMetric.from_config(**kwargs)`][evallm.metrics.base.BaseMetric.from_config])
    """

    model_config = ConfigDict(extra="forbid", frozen=True)


class BaseMetric(Generic[T], ABC):  # noqa: UP046
    """Abstract base class for all metrics, parameterized by a serializable value type.

    This generic class defines the shared interface implemented by all metrics. It
    includes convenience methods and properties for tree traversal, as well as an
    abstract `evaluate` method that must be implemented by all concrete metric
    classes.

    !!! info "Metric interface"
        - The `evaluate` method is expected to return a dictionary of the form
            `dict[str, T]`, where the key is this metric instance's
            [`result_key`][evallm.metrics.base.BaseMetric.result_key] and the value is
            of [`SerializableValueType`][evallm.metrics.types.SerializableValueType].
        - A metric's [`name`][evallm.metrics.base.BaseMetric.name] **identifies its
            class** in a [`Catalog`][evallm.metrics.registry.Catalog] and serves as the
            default `result_key` unless a custom result key is set on the instance. The
            [`result_key`][evallm.metrics.base.BaseMetric.result_key] on the other hand,
            is an **instance-level property** that determines the key under which this
            metric's value is stored in the evaluation results.
        - The (invariant) generic type parameter `T` declares the expected serializable
            value shape for a metric. At class-definition time, `BaseMetric` validates
            that `T` is compatible with
            [`SerializableValueType`][evallm.metrics.types.SerializableValueType], and
            `evaluate` is validated against `dict[str, T]`.
        - At runtime, calling [`result(value)`][evallm.metrics.base.BaseMetric.result]
            enforces that `value` matches `T` and returns `{self.result_key: value}`.
            Implementations that bypass `result(value)` are expected, by convention,
            to preserve the same contract.

    Keyword Args:
        **kwargs: The runtime configuration parameters for this metric instance, as
            a dictionary. The expected parameters depend on the metric class' defined
            [`config_schema`][evallm.metrics.base.BaseMetric.config_schema] and the
            provided parameters must strictly conform to it.

    Raises:
        ValidationError: If trying to instantiate a metric class with invalid runtime
            parameters (i.e. missing required parameters, extra unknown parameters, or
            parameters of the wrong type).
        ValueError: If trying to pass runtime parameters to a metric class that doesn't
                define a `config_schema`.
    """

    value_type: ClassVar[object]  # set in __init_subclass__
    """The metric class' parametrized type, binding the generic type parameter `T`."""
    config_schema: ClassVar[type[BaseModel]]
    """A schema defining the expected runtime configuration parameters for this metric.

    If the metric class does not explicitly define a config schema, this defaults to an
    empty pydantic model (see
    [`DefaultSchema`][evallm.metrics.base.DefaultSchema]) that takes no
    parameters.

    !!! info
        A class' config schema is always configured such that:

        - all declared fields must be provided unless they are declared as optional
        - no additional fields beyond those declared in the schema are allowed (see
            also [`pydantic.ConfigDict.extra`](https://docs.pydantic.dev/2.12/api/config/#pydantic.config.ConfigDict.extra))
        - the schema is frozen/immutable once defined (see also
            [`pydantic.ConfigDict.frozen`](https://docs.pydantic.dev/2.12/api/config/#pydantic.config.ConfigDict.frozen))
        - configuration parameters provided at runtime are strictly validated against
            the schema, raising a `ValidationError` if the provided parameters violate
            it (see also [pydantic Conversion Table](https://docs.pydantic.dev/2.12/concepts/conversion_table/#__tabbed_1_1))
    """
    _name: ClassVar[str]
    """The name of this metric class; used as registry identifier of subclasses."""
    _result_key: str | None = None
    """Optional key for this metric instance's value in the evaluation result
    dictionary.

    If `None`, the metric's [`name`][evallm.metrics.base.BaseMetric.name] is used as
    the default result key.
    """

    def __init_subclass__(cls, **kwargs):
        """Called when a class inherits from `BaseMetric`.

        Sets the default name and config schema of the metric class and validates
        & stores the parametrized value type of the metric class (i.e.
        the `T` in `BaseMetric[T]`) as the
        [`value_type`][evallm.metrics.base.BaseMetric.value_type] class variable
        and ensures that `evaluate` is implemented with the correct signature.

        Raises:
            TypeError: If the metric class does not implement the `evaluate` method
                with the correct signature.
            UnserializableValueTypeError: If the metric class's parametrized value type
                is not an instance of
                [`SerializableValueType`][evallm.metrics.types.SerializableValueType].
        """
        super().__init_subclass__(**kwargs)

        # Set default class name if not already set
        if getattr(cls, "_name", None) is None:
            cls._name = cls.__name__

        # Set default config schema if not already set
        if getattr(cls, "config_schema", None) is None:
            cls.config_schema = DefaultSchema

        # Set the metric's parametrized value type as a class variable
        for base in get_original_bases(cls):
            if get_origin(base) is BaseMetric:
                value_type = get_args(base)[0]
                if not is_serializable_type(value_type):
                    raise UnserializableValueTypeError(
                        f"Metric class '{cls.__name__}' has invalid value type annotation "
                        + f"{value_type!r}. Expected a SerializableValueType."
                    )
                cls.value_type = value_type
                break

        try:  # Ensure evaluate method is implemented & has correct signature
            _ = check_evaluate_signature(cls)
        except TypeError:
            raise

    def __init__(self, **kwargs: dict[str, Any]):
        """ """
        self.parent: BaseMetric[SerializableValueType] | None = None
        self._config: dict[str, Any] = {}  # default

        # raise error when trying to pass runtime params without a config schema
        if (schema := type(self).config_schema).model_json_schema().get(
            "properties", {}
        ) == {}:
            if kwargs:
                raise ValueError(f"{type(self)._name} takes no runtime parameters.")
        else:  # validate runtime params against config_schema
            try:
                model: BaseModel = schema.model_validate(kwargs)
                self._config = model.model_dump()
            except ValidationError:
                raise

    def __eq__(self, other: object) -> bool:
        """Check if this metric is equal to another metric.

        Two metrics are considered equal iff they are the same instance.

        Args:
            other (object): The other metric to compare against.

        Returns:
            bool: `True` if this metric is the same instance as the other metric,
                `False` otherwise.
        """
        return self is other

    def __iter__(self) -> Iterator[BaseMetric[SerializableValueType]]:
        """Iterate over the subtree rooted at this metric in breadth-first order.

        Yields:
            metric (BaseMetric): An iterator over the subtree rooted at this metric in
                breadth-first order, excluding this metric itself.
        """
        yield from self.descendants(include_self=False, order="bfs")

    def __contains__(self, metric: str | BaseMetric[SerializableValueType]) -> bool:
        """Check if a metric is contained in the subtree rooted at this metric.

        Note:
            Passing a metric **instance** checks for **identity** (i.e. whether exactly
            that instance is contained in this subtree), whereas passing a metric's
            **name** checks for **membership** (i.e. whether any metric with that name
            is contained in this subtree).

        Example:
            ```python
            leaf1: MetricA = MetricA()  # name='MetricA'
            leaf2: MetricA = MetricA()  # name='MetricA'
            composite: CompositeMetric = CompositeMetric()
            composite.add(leaf1)
            ```
            Identity-based membership check:
            ```pycon
            >>> leaf1 in composite
            True
            >>> leaf2 in composite
            False
            ```
            Name-based membership check:
            ```pycon
            >>> "MetricA" in composite
            True
            >>> "NonExistentMetric" in composite
            False
            ```

        Args:
            metric (str | BaseMetric): The name of the metric to check for, or the
                metric instance itself.

        Returns:
            bool: `True` if the given metric is a descendant of this metric
                (excluding this metric itself), `False` otherwise.
        """
        if isinstance(metric, str):
            return any(node.name == metric for node in self.descendants())
        return any(node is metric for node in self.descendants())

    def _make_repr(self, indent: int = 4, is_last: bool = True) -> str:
        """Helper method to create a string representation of the metric instance in a
        hierarchical format.

        Args:
            indent (int): The number of spaces to use for each level of indentation.
            is_last (bool): Whether this metric is the last child of its parent. Used to
                determine the appropriate tree branch characters for the representation.

        Returns:
            representation (str): A string representation of the metric instance,
                including its name and class, formatted with indentation based on the
                level in the hierarchy.
        """

        def attrs() -> str:

            def stringify_value_type(value_type: object) -> str:  # pragma: no cover
                origin = get_origin(value_type)
                args = get_args(value_type)

                if origin is list and len(args) == 1:
                    return f"list[{stringify_value_type(args[0])}]"
                if origin is tuple and args:
                    return (
                        f"tuple[{', '.join(stringify_value_type(arg) for arg in args)}]"
                    )
                if origin is dict and len(args) == 2:
                    key_type, val_type = args
                    return f"dict[{stringify_value_type(key_type)}, {stringify_value_type(val_type)}]"
                if value_type is None or value_type is type(None):
                    return "None"
                if isinstance(value_type, type):
                    return value_type.__name__
                return str(value_type)

            return " ".join(
                [
                    f"{self.__class__.__name__} at {hex(id(self)).upper()}:",
                    f"value_type={stringify_value_type(self.value_type)}",
                ]
            )

        def prefix(is_last: bool) -> str:
            prefix = " " * indent * self.depth
            if self.is_root:
                return prefix
            return prefix + ("└─ " if is_last else "├─ ")

        return f"{prefix(is_last=is_last)}'{self.name}' ({attrs()})"

    def _repr_tree(self) -> str:
        lines: list[str] = []
        lines.append(self._make_repr(is_last=True))  # root line

        if self.is_composite:
            children = list(self.children)
            for i, child in enumerate(children):
                is_last = i == len(children) - 1
                lines.extend(child._repr_tree_lines(is_last=is_last))
        return "\n".join(lines)

    def _repr_tree_lines(self, *, is_last: bool) -> list[str]:
        # render this node as a child of someone else
        lines = [self._make_repr(is_last=is_last)]

        if self.is_composite:
            children = list(self.children)
            for i, child in enumerate(children):
                child_is_last = i == len(children) - 1
                lines.extend(child._repr_tree_lines(is_last=child_is_last))
        return lines

    def __repr__(self) -> str:
        """String representation of the subtree, rooted in this metric instance.

        Returns:
            representation (str): A string representation of the metric instance,
                including its name and class.
        """
        return self._repr_tree()

    def result(self, value: T) -> dict[str, T]:  # pragma: no cover
        """Helper method for wrapping the value of a metric's evaluation result in a
        dictionary to conform to the expected return type of `evaluate`.

        Example:
            ```python
            class CustomMetric(BaseMetric[list[int]]):
                def evaluate(self, text: str) -> dict[str, list[int]]:
                    # Perform some evaluation logic and compute the metric's value
                    res: list[int] = [len(text.split())]

                    # Wrap the result in a dictionary keyed by the metric's result_key
                    return self.result(res)
            ```

        Args:
            value (T): The result of computing the metric, where `T` is the metric's
                parametrized value type.

        Returns:
            result (dict[str, T]): A dictionary containing the given value, keyed by
                the metric's [`result_key`][..result_key].

        Raises:
            TypeError: If the given value is not of the expected type `T` or is not
                JSON-serializable.
        """
        if not matches_annotation(value, self.value_type):
            raise TypeError(
                f"Value {value!r} of type {type(value)!r} is not of expected type "
                + f"{self.value_type!r} for metric '{self.name}'."
            )
        return {self.result_key: value}

    @abstractmethod
    def evaluate(self, text: str) -> dict[str, T]:
        """Evaluate the metric on the given text.

        Args:
            text (str): The text to evaluate the metric on.

        Returns:
            value (dict[str, T]): The evaluation result keyed by the metric's
                [`result_key`][..result_key], where `T` is the metric's parametrized
                value type.
        """
        raise NotImplementedError(  # pragma: no cover
            "Subclasses must implement evaluate() method."
        )

    #########################
    # Read-write properties #
    #########################
    @property
    def parent(self) -> CompositeMetric | None:
        """Get the parent of this metric.

        Returns:
            CompositeMetric | None: The parent metric, or `None` if
                this metric has no parent, i.e. is a leaf node.
        """
        return self._parent

    @parent.setter
    def parent(self, parent: CompositeMetric | None) -> None:
        """Set the parent of this metric.

        This is used to build a tree of metrics, where each metric can have a parent
        and multiple children.

        Args:
            parent (CompositeMetric | None): The parent metric to set.
                Can be `None` to unset the parent.
        """
        self._parent: CompositeMetric | None = parent

    @property
    def result_key(self) -> str:
        """The key to use for this metric's value in the evaluation result dictionary.

        By default, this is the same as the metric's [`name`][..name], but it can be
        overridden on a per-instance basis. This is useful when multiple instances of
        the same metric class are used in a single evaluation request with different
        runtime parameters and therefore need distinct keys in the output.

        Returns:
            result_key (str): The key of this metric instance's value in the evaluation
                result dictionary. Defaults to the metric's [`name`][..name].
        """
        if self._result_key is None:
            return self.name
        return self._result_key

    @result_key.setter
    def result_key(self, key: str | None) -> None:
        """Set the key to use for this metric's value in the evaluation result
        dictionary.

        Args:
            key (str | None): A non-empty key to use or `None` to use the metric's
                [`name`][..name] as the result key.

        Raises:
            TypeError: If the provided key is not a string or `None`.
            ValueError: If the provided key is an empty string or contains only
                whitespace.
        """
        if not isinstance(key, (str, type(None))):
            raise TypeError("result_key must be a string or None.")
        if key is not None and key.strip() == "":
            raise ValueError(
                "result_key cannot be an empty string or contain only whitespace."
            )
        self._result_key = key

    #########################
    # Read-only properties  #
    #########################
    @property
    def name(self) -> str:
        """The name of this metric's class.

        This is a class-, not an instance-level property that's used for identifying
        this metric in the catalog. It also serves as the default
        [`result_key`][..result_key] for metric instances unless explicitly overridden.

        !!! note "Metric class registration and name updates"
            When a metric class is registered in the catalog under a custom name using
            `@register_metric(name=...)`, the name-property of any pre-existing
            instances of that class will be updated to reflect the new name in the
            catalog.

        Returns:
            name (str): The name of this metric. Defaults to the class name if not
                overwritten by the `@register_metric` decorator at definition time.
        """
        return self.__class__._name

    @property
    def is_composite(self) -> bool:
        """Check if this metric is a composite metric (i.e. can have child metrics).

        Returns:
            bool: `True` if this metric is a composite metric, `False` otherwise.
        """
        return hasattr(self, "_children")

    @property
    def is_root(self) -> bool:
        """Check if this metric is a root metric (i.e. has no parent).

        Returns:
            bool: `True` if this metric has no parent, `False` otherwise.
        """
        return self._parent is None

    @property
    def root(self) -> BaseMetric[SerializableValueType]:
        """Get the root metric of the tree this metric belongs to.

        Returns:
            BaseMetric: The root metric of the tree this metric belongs to.
        """
        node: BaseMetric[SerializableValueType] = cast(
            BaseMetric[SerializableValueType], self
        )
        while node.parent is not None:
            node = cast(BaseMetric[SerializableValueType], node.parent)
        return node

    @property
    def depth(self) -> int:
        """Get the depth of this metric in the tree (i.e. the number of edges from this
        metric to the root).

        Returns:
            depth (int): The depth of this metric in the tree. Root metrics have
                depth 0, their children have depth 1, and so on.
        """
        parent: CompositeMetric | None = self.parent
        if parent is None:
            return 0
        return 1 + parent.depth

    @property
    def children(self) -> Iterator[BaseMetric[SerializableValueType]]:
        """Get an iterator over the children of this metric.

        Yields:
            children (BaseMetric): An iterator over the child metrics (i.e. direct
                descendants) of this metric. If this metric is a leaf node, this
                will be an empty iterator.
        """
        if not self.is_composite:
            return
        yield from self._children  # type: ignore[attr-defined]

    def descendants(
        self,
        *,
        include_self: bool = False,
        order: Literal["dfs", "bfs"] = "dfs",
    ) -> Iterator[BaseMetric[SerializableValueType]]:
        """Yield nodes in the subtree rooted at self (cycle-safe).

        Args:
            include_self (bool): Whether to include `self` in the yielded nodes.
            order ("dfs" | "bfs"): The order to traverse the tree. "dfs" for depth-first
                search, "bfs" for breadth-first search.

        Yields:
            descendants (Iterator[BaseMetric]): An iterator over the descendant nodes of this metric,
                in the specified traversal order.
        """
        visited: set[int] = set()

        # Populate the frontier with the initial node(s) based on traversal order
        frontier: deque[BaseMetric[SerializableValueType]] = deque()
        if include_self:
            frontier.append(cast(BaseMetric[SerializableValueType], self))
        else:  # keep order of children consistent across DFS/BFS
            if order == "dfs":
                frontier.extend(reversed(list(self.children)))
            else:
                frontier.extend(list(self.children))

        while frontier:
            # Get the next node to visit, depending on traversal order
            node: BaseMetric[SerializableValueType] = (
                frontier.pop() if order == "dfs" else frontier.popleft()
            )

            # Skip visited nodes to prevent infinite loops
            if (node_id := id(node)) in visited:
                continue
            visited.add(node_id)

            yield node

            # Expand neighbors
            if order == "dfs":
                frontier.extend(  # reverse for correct DFS order
                    reversed(list(node.children))
                )
            else:
                frontier.extend(list(node.children))

    @property
    def config(self) -> dict[str, Any]:
        """Get the runtime configuration of this metric instance as a dictionary.

        Returns:
            config (dict[str, Any]): A dictionary containing the runtime configuration
                of this metric instance, validated against the metric's schema.
        """
        return self._config

    @classmethod
    def from_config(
        cls, *, json_payload: str | None = None, **kwargs: Any
    ) -> BaseMetric[T]:
        """Factory method to create a metric instance from kwargs or JSON payload.

        !!! info
            This method can be used to create any metric instance, so long as the
            provided config parameters conform to the metric's schema. If the metric
            class doesn't define a config schema (the default), calling this method
            without kwargs is equivalent to calling the metric class constructor
            directly.
            The expected parameters depend on the metric's declared config_schema
            (see also [`@config_schema` decorator][evallm.metrics.config.config_schema])
            and the provided parameters must strictly conform to it.

        Args:
            json_payload (str | None): A JSON string containing the runtime parameters
                for this metric instance.

        Keyword Args:
            **kwargs (Any): The runtime parameters to use for this metric instance.

        Returns:
            metric (BaseMetric[T]): An instance of this metric class created using the
                provided configuration parameters.

        Raises:
            TypeError: If `json_payload` is not a string or `None`.
            ValueError: If

                - trying to pass runtime parameters to a metric class that doesn't
                    define a `config_schema`, or
                - trying to instantiate a metric class with invalid runtime parameters
                    (i.e. missing required parameters, extra unknown parameters, or
                    parameters of the wrong type), or
                - passing an invalid JSON payload, or
                - passing both `json_payload` and `kwargs` at the same time.
        """
        runtime_params: dict[str, Any]
        if json_payload is not None:
            if kwargs:
                raise ValueError("Cannot mix json_payload with kwargs.")
            if not isinstance(json_payload, str):
                raise TypeError("json_payload must be a string or None.")
            try:
                runtime_params = orjson.loads(json_payload)
            except orjson.JSONDecodeError:
                raise ValueError(
                    "Invalid json_payload: not a valid JSON string."
                ) from None
        else:
            runtime_params = kwargs

        return cls(**runtime_params)


__all__ = [
    "BaseMetric",
    "DefaultSchema",
]
