#!/usr/bin/env python3

"""Global registry for automatic metrics discovery

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

from __future__ import annotations

from collections.abc import Callable
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Protocol,
    overload,
    runtime_checkable,
)

from .base import BaseMetric
from .types import SerializableValueType

if TYPE_CHECKING:
    pass


####################
#     Catalog      #
####################
@runtime_checkable
class Catalog(Protocol):
    """Interface for a global catalog that keeps track of distinct available metrics.

    Defines a `register(metric: BaseMetric[SerializableValueType])` method that serves
    as a hook for the `@register_metric` decorator, enabling automatic registration of
    `BaseMetric[SerializableValueType]`-subclasses in the catalog at definition time.
    Also provides convenience methods for looking up metric classes by name and
    renaming registered metrics.
    """

    @classmethod
    def available(cls) -> list[str]: ...

    @classmethod
    def get_metric(cls, identifier: str) -> type[BaseMetric[SerializableValueType]]: ...

    @classmethod
    def is_registered(
        cls,
        metric: (
            BaseMetric[SerializableValueType] | type[BaseMetric[SerializableValueType]]
        ),
    ) -> bool: ...

    @classmethod
    def register(cls, metric: type[BaseMetric[SerializableValueType]]) -> None: ...

    @classmethod
    def rename_metric(cls, old_name: str, new_name: str) -> None: ...


class MetricsCatalog:
    """A global catalog that keeps track of distinct available metrics.

    Implements the `Catalog` protocol by maintaining an internal registry of metric
    names mapped to their corresponding `BaseMetric[SerializableValueType]` subclasses.
    Provides methods for looking up metric classes by name and renaming registered
    metrics.

    Important:
        In order for metrics to be registered in the catalog, the module in which they
        are defined and decorated has to be imported.
    """

    _registry: ClassVar[dict[str, type[BaseMetric[SerializableValueType]]]] = {}

    @classmethod
    def get_metric(cls, name: str) -> type[BaseMetric[SerializableValueType]]:
        """Get the metric class corresponding to the given name.

        Args:
            name (str): The name of the metric to retrieve.

        Returns:
            metric (type[BaseMetric[SerializableValueType]]): The metric class
                corresponding to the given name.

        Raises:
            KeyError: If no metric with the given name is registered in the
                catalog.
        """
        try:
            return cls._registry[name]
        except KeyError:
            raise KeyError(
                f"Metric with name '{name}' is not registered in the catalog."
            ) from None

    @classmethod
    def available(cls) -> list[str]:
        """Get a list of registered metrics' names.

        Returns:
            names (list[str]): A list of the names of all metrics currently registered
                in the catalog.
        """
        return list(cls._registry.keys())

    @classmethod
    def is_registered(
        cls,
        metric: (
            BaseMetric[SerializableValueType] | type[BaseMetric[SerializableValueType]]
        ),
    ) -> bool:
        """Check if a given metric is registered in the catalog.

        Args:
            metric (BaseMetric[SerializableValueType] | type[BaseMetric[SerializableValueType]]): A metric to check for. Can be
                either a `BaseMetric[SerializableValueType]` instance or a `BaseMetric[SerializableValueType]` subclass.

        Returns:
            bool: `True` if the given metric is registered in the catalog, `False`
                otherwise.

        Raises:
            TypeError: If the provided `metric` is not an instance or subclass
                of `BaseMetric[SerializableValueType]`.
        """
        if not (
            isinstance(metric, BaseMetric)
            or (isinstance(metric, type) and issubclass(metric, BaseMetric))
        ):
            raise TypeError(
                f"Expected a `{type(BaseMetric[SerializableValueType]).__name__}` instance or subclass "
                + f"but got {type(metric).__name__}."
            )
        metric_cls = metric if isinstance(metric, type) else metric.__class__
        if (name := getattr(metric_cls, "_name", None)) is None:  # pragma: no cover
            return False  # shouldn't happen
        return any(m_cls._name == name for m_cls in cls._registry.values())

    @classmethod
    def register(cls, metric: type[BaseMetric[SerializableValueType]]) -> None:
        """Registers a metric in the catalog.

        Args:
            metric (type[BaseMetric[SerializableValueType]]): The `BaseMetric[SerializableValueType]` subclass to register.

        Raises:
            TypeError: If trying to register a class that is not a subclass of
                `BaseMetric[SerializableValueType]`.
            ValueError: If a metric with the same name is already registered.
        """
        try:
            if cls.is_registered(metric):
                raise ValueError(
                    f"Metric '{metric._name}' is already registered in the catalog."
                )
        except TypeError as e:  # pragma: no cover
            raise e from None
        cls._registry[metric._name] = metric

    @classmethod
    def rename_metric(cls, old_name: str, new_name: str) -> None:
        """Renames a metric and updates its references in the catalog.

        Args:
            old_name (str): The current name of the metric to rename.
            new_name (str): The new name to assign to the metric.

        Raises:
            KeyError: If no metric with `old_name` is currently registered.
            ValueError: If

                - another metric with `new_name` is already registered in the catalog, or
                - `new_name` is an empty string.
        """
        if old_name not in cls._registry:
            raise KeyError(f"Metric '{old_name}' is not registered in the catalog.")
        if new_name in cls._registry:
            raise ValueError(
                f"Metric '{new_name}' is already registered in the catalog."
            )
        if not isinstance(new_name, str) or new_name == "":
            raise ValueError("Metric name must be a non-empty string.")

        # Rename the metric class
        metric_cls: type[BaseMetric[SerializableValueType]] = cls._registry.pop(
            old_name
        )
        metric_cls._name = new_name

        # update the reference in the registry
        cls._registry[new_name] = metric_cls


####################
# @register_metric #
####################
@overload
def register_metric(  # pragma: no cover
    cls: type[BaseMetric[SerializableValueType]],
    /,
) -> type[BaseMetric[SerializableValueType]]: ...


@overload
def register_metric(  # pragma: no cover
    cls: None = None,
    /,
    *,
    catalog: type[Catalog] = MetricsCatalog,
    name: str | None = None,
) -> Callable[
    [type[BaseMetric[SerializableValueType]]],
    type[BaseMetric[SerializableValueType]],
]: ...


def register_metric(
    cls: type[BaseMetric[SerializableValueType]] | None = None,
    /,
    *,
    catalog: type[Catalog] = MetricsCatalog,
    name: str | None = None,
) -> (
    type[BaseMetric[SerializableValueType]]
    | Callable[
        [type[BaseMetric[SerializableValueType]]],
        type[BaseMetric[SerializableValueType]],
    ]
):
    """Decorator for registering a metric class in a global `Catalog` class.

    Decorated metric classes are automatically registered at *import time* and thus the
    module in which a given metric is defined has to be imported in order for it to
    become available for lookup in the catalog.

    Tip:
        You can use a custom catalog in which to register your metrics (e.g. because
        you want to maintain separate catalogs for different types of metrics) by
        implementing the `Catalog` protocol.

    Args:
        catalog (type[Catalog]): An optional `Catalog` class to register the metric in.
            Uses `MetricsCatalog` by default to globally register metrics in a single
            catalog.
        name (str | None): An optional name to register the metric under. Must be unique
            among registered metric names in the catalog, and cannot be an empty string
            if provided. If not provided, the metric will be registered under its class
            name.

    Returns:
        decorator
            (type[BaseMetric[SerializableValueType]] | Callable[
                [type[BaseMetric[SerializableValueType]]],
                type[BaseMetric[SerializableValueType]]]): The decorated metric class (
                    `@register_metric`) or a decorator function (
                    `@register_metric(...)`) that validation and registration.

    Raises:
        TypeError: If

            - the value type `T` of the decorated class'
                `evaluate(...) -> dict[str, T]` method can't be determined, or
            - the decorated class is not a subclass of `BaseMetric[SerializableValueType]`.
        UnserializableValueTypeError: If the value type `T` of the decorated class'
            `evaluate(...) -> dict[str, T]` method is not a subclass of
            `SerializableValueType`.
        ValueError: If

            - an empty string is provided as the `name` argument, or
            - a metric with the same name is already registered in the catalog.
    """  # noqa: DOC502

    def decorate(
        cls: type[BaseMetric[SerializableValueType]],
    ) -> type[BaseMetric[SerializableValueType]]:
        """Decorate a metric class and register it in the provided `Catalog` instance.

        Args:
            cls (type[BaseMetric[SerializableValueType]]): The metric class to
                register. Must be a subclass of `BaseMetric[SerializableValueType]`.

        Returns:
            cls (type[BaseMetric[SerializableValueType]]): The decorated metric class
                after having registered it.
        """
        cls._name = name or cls._name  # optionally set a custom name
        catalog.register(metric=cls)
        return cls

    # validate the provided name before attempting to register the metric class
    if name is not None and (not isinstance(name, str) or name == ""):
        raise ValueError("Metric name must be a non-empty string.")

    if cls is not None:  # no args: `@register_metric`
        return decorate(cls)
    return decorate  # with args: `@register_metric(...)`


__all__ = [
    "Catalog",
    "MetricsCatalog",
    "register_metric",
]
