#!/usr/bin/env python3

"""`CompositeMetric` — tree container for other metrics

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

from __future__ import annotations

from typing import cast

from .base import BaseMetric
from .types import SerializableValueType


class CompositeMetric(BaseMetric[dict[str, SerializableValueType]]):
    """A container for other metrics following the `BaseMetric` interface.

    Note:
        Implements the *Composite* design pattern by allowing a metric to contain other
        `BaseMetric` & `CompositeMetric` instances so that calling `evaluate` on a
        composite metric will recursively call `evaluate` on all of its children.
    """

    def __init__(self) -> None:
        super().__init__()
        self._children: list[BaseMetric[SerializableValueType]] = []

    #############
    # Composite #
    #############
    def add(self, metric: BaseMetric[SerializableValueType], *, move: bool = False):
        """Add a child metric to this composite metric instance.

        Note:
            A composite metric can contain any number of child metrics, which can be
            simple metrics or other composite. The only restriction is that the *same*
            *instance* of a metric can only be added to exactly one composite metric at
            any given time so as to ensure a well-formed tree structure.

            Note however that *different instances* of the *same class* of metric can be
            added arbitrarily often to the same or to different composite metrics.

        Args:
            metric (BaseMetric[SerializableValueType]): The child metric to add (can
                be simple or composite). The parent of the child will be set to this
                composite metric. If this composite is already the metric's parent,
                this is a no-op.
            move (bool): Whether to remove the `metric` from its current parent (if
                any) before adding it to this composite. If `False`, trying to add a
                metric that already has a parent will raise a `ValueError`.

        Raises:
            ValueError: If trying to

                - add self as a child of itself, or
                - add a metric instance that is already contained in this
                    composite metric, or
                - add a metric instance that is already contained in another composite
                    metric and `move=False`.
        """
        if metric is self:
            raise ValueError("Cannot add self as a child of itself.")
        if metric in self:
            if metric.parent is self:
                return
            raise ValueError(
                f"{metric!r} is already a descendant of this composite metric."
            )
        if metric.parent is not None:
            if not move:
                raise ValueError(
                    f"{metric!r} is already a descendant of {metric.parent!r}\n"
                    + "Set move=True to automatically remove it from its current "
                    + "parent and add it to this composite metric."
                )
            # remove from current parent before adding to this composite
            metric.parent.remove(metric)

        self._children.append(metric)
        metric.parent = self

    def _remove_child(
        self, metric: BaseMetric[SerializableValueType]
    ) -> BaseMetric[SerializableValueType]:
        """Recursively search for the given metric's parent in this composite's subtree
        and remove the metric from its children.

        After this operation, the removed metric's parent will be unset (set to `None`).

        Args:
            metric (BaseMetric[SerializableValueType]): The child metric to remove.

        Returns:
            removed_metric (BaseMetric[SerializableValueType]): The metric that was
                removed.

        Raises:
            ValueError: If the specified metric is not a descendant of this composite.
        """
        for i, child in enumerate(self._children):
            if child is metric:
                removed_metric: BaseMetric[SerializableValueType] = self._children.pop(
                    i
                )
                removed_metric.parent = None
                return removed_metric
            elif child.is_composite:  # recursively search in composite children
                try:
                    return cast(CompositeMetric, child)._remove_child(metric)
                except ValueError:  # continue searching in other branches
                    pass
        raise ValueError(f"{metric!r} is not a descendant of {self!r}")

    def remove(
        self, metric: BaseMetric[SerializableValueType]
    ) -> BaseMetric[SerializableValueType]:
        """Remove a contained metric from this composite.

        Important:
            When removing a composite metric, all of its descendants are also removed.

        Args:
            metric (BaseMetric[SerializableValueType]): The child metric to remove.
                The parent of the child will be unset (set to `None`).

        Returns:
            removed_metric (BaseMetric[SerializableValueType]): The metric that was
                removed.

        Raises:
            ValueError: If

                - trying to remove this composite from itself, or
                - the specified metric is not a descendant of this composite.
        """
        if metric is self:
            raise ValueError("Cannot remove self from itself.")
        return self._remove_child(metric)

    def evaluate(self, text: str) -> dict[str, dict[str, SerializableValueType]]:
        """Evaluate all child metrics on the given text.

        This method calls `evaluate` on all child metrics and return the results as
        a nested dictionary.

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

        Returns:
            result (dict[str, dict[str, SerializableValueType]]): A dictionary
                containing the results of all child metric evaluations, where each
                child's result contributes entries keyed by that child's
                [`result_key`][evallm.metrics.base.BaseMetric.result_key]. The combined
                child results are then wrapped in a dictionary under this composite's
                own name.
        """
        res: dict[str, SerializableValueType] = {}
        for c in self._children:
            res.update(c.evaluate(text))
        return {self.name: res}


__all__ = [
    "CompositeMetric",
]
