#!/usr/bin/env python3

"""Types and serialization helper functions

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

from __future__ import annotations

import datetime as dt
from collections.abc import Mapping, Sequence
from types import GenericAlias, ModuleType
from typing import (
    Any,
    ClassVar,
    Protocol,
    get_args,
    get_origin,
    get_type_hints,
    runtime_checkable,
)
from uuid import UUID

import orjson
from pydantic import BaseModel


@runtime_checkable
class MetricClassProto(Protocol):
    _name: ClassVar[str]
    value_type: ClassVar[object]
    config_schema: ClassVar[type[BaseModel]]


# Atomic values orjson can serialize without `default=...`
type OrjsonLeaf = (
    str
    | int
    | float
    | bool
    | dt.datetime
    | dt.date
    | dt.time
    | UUID
    | orjson.Fragment
    | None
)
ORJSON_LEAVES: set[type] = {
    str,
    int,
    float,
    bool,
    dt.datetime,
    dt.date,
    dt.time,
    UUID,
    orjson.Fragment,
    type(None),
}
"""Atomic serializable types for runtime typechecking."""

# Recursive container types
type SerializableValueType = (
    OrjsonLeaf
    # Note: while tuples are natively supported by orjson, subclasses of tuple
    # (e.g. namedtuple) are not
    | tuple["SerializableValueType", ...]
    | list["SerializableValueType"]
    | dict[str, "SerializableValueType"]
)
"""Allowed parameterization types for the `T` in `BaseMetric[T]`.

The `SerializableValueType` type contains all types that can be used as the value
type `T` in subclasses' implementations of the abstract
`BaseMetric.evaluate(text: str) -> dict[str, T]` method.

Supported types include:
    - `str`
    - `int`
    - `float`
    - `bool`
    - `dt.datetime`
    - `dt.date`
    - `dt.time`
    - `UUID`
    - `orjson.Fragment`

optionally wrapped in `tuple[...]`, `list[...]` or `dict[str, ...]`.

!!! warning
    While `tuple` is natively serialized, *subclasses of* `tuple` *(e.g.* `namedtuple`*) are not*.

!!! info "See also"
    [:material-github: orjson: Serialize](https://github.com/ijl/orjson?tab=readme-ov-file#serialize)
"""


class UnserializableValueTypeError(TypeError):
    """Exception raised when parametrization invariants of `BaseMetric[T]` are violated."""

    pass


def check_evaluate_signature(metric_cls: type[MetricClassProto]) -> type:
    """Helper function to ensure that a metric class has an `evaluate` method with an
    annotated signature that conforms to `BaseMetric`'s invariants.

    If `metric_cls.evaluate(text: str)` doesn't have a return annotation, it will be
    added using the class'
    [`metric_cls.value_type`][evallm.metrics.base.BaseMetric.value_type].

    !!! info "Expected signature"
        Every metric's `evaluate(text: str) -> dict[str, T]` method returns a
        dictionary whose values must conform to the metric's declared `T`. For leaf
        metrics, the expected convention is to return a single key-value pair using
        the metric instance's
        [`result_key`][evallm.metrics.base.BaseMetric.result_key] (which defaults to
        the metric's `name`). The generic type `T` is parametrized by each concrete
        metric class such that `T` is bound to any JSON-serializable type (see
        [`SerializableValueType`][evallm.metrics.types.SerializableValueType]).

    Args:
        metric_cls (type): The metric class to check.

    Returns:
        metric_cls (type): The same metric class (potentially with added return
            annotation for `metric_cls.evaluate`), if it passes all checks.

    Raises:
        NotImplementedError: If `metric_cls` doesn't implement an `evaluate` method.
        TypeError: If...
            - `metric_cls` is not a subclass of `BaseMetric`, or
            - `metric_cls` doesn't have an `evaluate` method, or
            - the method doesn't have a return type annotation, or
            - if the return type is not of shape `dict[str, T]`.
    """

    def make_return_annotation(value_type: object) -> Any:
        return GenericAlias(
            dict, (str, value_type)
        )  # equivalent to dict[str, value_type]

    from importlib import import_module

    base: ModuleType = import_module("evallm.metrics.base")  # avoid circular import
    BaseMetric: type = getattr(base, "BaseMetric")  # noqa: B009

    if not issubclass(metric_cls, BaseMetric):
        raise TypeError(
            f"metric_cls of type {metric_cls!r} is not a subclass of {BaseMetric!r}."
        )

    if (evaluate := getattr(metric_cls, "evaluate", None)) is None:  # pragma: no cover
        # shouldn't happen; primarily defensive
        raise NotImplementedError(
            f"Metric class {metric_cls.__name__} has no `evaluate` method."
        )
    elif getattr(evaluate, "__isabstractmethod__", False):
        raise NotImplementedError(
            f"Metric class {metric_cls.__name__} doesn't implement an `evaluate` method."
        )

    if (return_ann := get_type_hints(evaluate).get("return")) is None:
        # add a return annotation if missing, using the metric's `value_type` property
        evaluate.__annotations__["return"] = make_return_annotation(
            metric_cls.value_type
        )
    else:  # ensure return annotation is of shape `dict[str, metric_cls.value_type]`
        if get_origin(return_ann) is not dict:
            raise TypeError(
                f"Return type annotation {return_ann!r} of `evaluate` method of metric class "
                + f"{metric_cls.__name__} is not of the form `dict[str, T]` for some type `T`."
            )
        key_ann, value_ann = get_args(return_ann)
        if key_ann is not str:
            raise TypeError(
                f"Key type in return type annotation of {metric_cls!r}.evaluate must "
                + f"be {type(str)}, but got {key_ann}."
            )
        if value_ann != metric_cls.value_type:
            raise TypeError(
                "Mismatch between return type annotation of `evaluate` and declared "
                + f"value type `{metric_cls.value_type!r}` for metric class "
                + f"{metric_cls.__name__!r}."
            )

    return metric_cls


def is_serializable_type(tp: object) -> bool:
    """Helper function to recursively verify a given **type annotation**.

    Args:
        tp (object): The object to check for JSON-serializability.

    Returns:
        bool: `True` if `tp` is JSON-serializable, and `False` otherwise.
    """

    def is_namedtuple_subclass(t: type) -> bool:
        return isinstance(t, type) and issubclass(t, tuple) and hasattr(t, "_fields")

    if tp is SerializableValueType:
        return True
    origin: Any | None = get_origin(tp)

    if origin is None:
        if tp is None:
            tp = type(None)  # treat explicit runtime use of None as a type annotation

        if isinstance(tp, type) and is_namedtuple_subclass(tp):  # pragma: no cover
            return False

        if tp is tuple:
            return False  # must be tuple[T, ...] | tuple[T1, T2, ...]

        return tp in ORJSON_LEAVES

    if origin is list:
        (arg,) = get_args(tp)
        return is_serializable_type(arg)

    if origin is dict:
        k, v = get_args(tp)
        return k is str and is_serializable_type(v)

    if origin is tuple:
        args: tuple[Any, ...] = get_args(tp)

        # tuple[T, ...]
        if len(args) == 2 and args[1] is Ellipsis:
            return is_serializable_type(args[0])

        # tuple[T1, T2, ...] -> all declared types must be uniform and serializable
        first = args[0]
        return all(a == first for a in args) and is_serializable_type(first)

    return False


def is_serializable_value(v: object) -> bool:
    """Helper function to recursively verify a given **runtime value**.

    Args:
        v (object): The value to check for JSON-serializability.

    Returns:
        is_serializable (bool): `True` if `v` is JSON-serializable, `False` otherwise.
    """

    def is_namedtuple_instance(obj: object) -> bool:
        return isinstance(obj, tuple) and hasattr(obj, "_fields")

    if isinstance(v, tuple(ORJSON_LEAVES)):
        return True

    if isinstance(v, Mapping):
        return all(
            isinstance(k, str) and is_serializable_value(val) for k, val in v.items()
        )

    if isinstance(v, Sequence) and not isinstance(v, (str, bytes, bytearray)):
        if is_namedtuple_instance(v):
            return False
        return all(is_serializable_value(x) for x in v)

    return False


def matches_annotation(val: Any, ann: Any) -> bool:
    """Helper function to check whether a given runtime value matches a given type
    annotation.

    Args:
        val (Any): The runtime value to check.
        ann (Any): The type annotation to check against.

    Returns:
        matches (bool): `True` if `val` matches the type annotation `ann`, and `False`
            otherwise.
    """
    # since SerializableValueType is a recursive type, we need to check the value
    # itself rather than the annotation
    if ann is SerializableValueType:
        return is_serializable_value(val)

    # otherwise, we unpack the annotation's structure and check the value against it
    # recursively
    origin = get_origin(ann)
    if origin is None:
        if ann is None:
            return val is None
        return isinstance(val, ann)

    if origin is list:
        (item_ann,) = get_args(ann)
        return isinstance(val, list) and all(
            matches_annotation(x, item_ann) for x in val
        )

    if origin is dict:
        k_ann, v_ann = get_args(ann)
        return (
            isinstance(val, dict)
            and k_ann is str
            and all(
                isinstance(k, str) and matches_annotation(v, v_ann)
                for k, v in val.items()
            )
        )

    if origin is tuple:
        args = get_args(ann)

        if not isinstance(val, tuple):
            return False

        # tuple[T, ...]  (homogeneous, any length)
        if len(args) == 2 and args[1] is Ellipsis:
            item_ann = args[0]
            return all(matches_annotation(x, item_ann) for x in val)

        # tuple[T1, T2, ...]  (fixed length)
        if len(val) != len(args):
            return False
        return all(matches_annotation(x, t) for x, t in zip(val, args, strict=True))

    return False


__all__ = [
    "SerializableValueType",
    "UnserializableValueTypeError",
]
