#!/usr/bin/env python3

"""Configuration schema & runtime parameters for metrics

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

from __future__ import annotations

from collections.abc import Callable
from types import UnionType
from typing import (
    TYPE_CHECKING,
    Any,
    NotRequired,
    Protocol,
    Required,
    TypeGuard,
    Union,
    cast,
    get_args,
    get_origin,
    get_type_hints,
)

from pydantic import BaseModel, ConfigDict, TypeAdapter, create_model

if TYPE_CHECKING:
    from .base import BaseMetric


####################
#  @config_schema  #
####################
if TYPE_CHECKING:

    class TypedDictType(Protocol):
        __annotations__: dict[str, Any]
        __total__: bool
        __required_keys__: frozenset[str]
        __optional_keys__: frozenset[str]


type PlainDictSchema = dict[str, object]
"""A plain dictionary mapping configuration parameters to their types.

All parameters defined in a plain dictionary schema are considered required.

!!! info
    The dictionary's values must be **type objects** that serve as runtime annotations.
    Make sure you're passing the *expected types* rather than the parameter values
    themselves.
"""
type TypedDictSchema = type[TypedDictType]
"""A `TypedDict` subclass defining configuration parameters and their types."""
type PydanticSchema = type[BaseModel]
"""A [:simple-pydantic: `pydantic.BaseModel`](https://docs.pydantic.dev/2.12/api/base_model/) subclass.
defining configuration parameters, their types, and optionally validation logic and
default values."""


class InvalidConfigSchema(TypeError):
    """Exception raised when an invalid schema is provided to the `@config_schema`
    decorator."""


def _is_typeddict_type(obj: object) -> TypeGuard[TypedDictSchema]:
    """Helper to heuristically check if an object is a `TypedDict` subclass at runtime.

    Args:
        obj (object): The object to check.

    Returns:
        TypeGuard[TypedDictSchema]: `True` if the object is a `TypedDict` subclass,
            `False` otherwise. (Note that this is a heuristic check and may not be
            100% reliable in all cases.)
    """
    # TypedDict subclasses carry these markers; instances are plain dicts
    return (
        isinstance(obj, type)
        and hasattr(obj, "__annotations__")
        and hasattr(obj, "__total__")
        and hasattr(obj, "__required_keys__")
        and hasattr(obj, "__optional_keys__")
    )


def _is_annotation_like(annotation: object) -> bool:
    """Check if `annotation` looks like a valid type
    [*annotation expression*](https://typing.python.org/en/latest/spec/annotations.html#annotation-expression).

    Accepts:
    - concrete types like `int`, `str`, `bool`
    - parametrized typing forms like `list[str]`, `dict[str, int]`, `float | None`
    - nested plain-dict schemas, validated elsewhere

    Args:
        annotation (object): The annotation to check.

    Returns:
        bool: `True` if the annotation looks like a valid type annotation, `False`
            otherwise.
    """
    if annotation is None or annotation is ...:
        return False

    if isinstance(annotation, dict):
        return False

    if isinstance(annotation, type):
        return True

    # value-like runtime objects won't have an origin
    return get_origin(annotation) is not None


def _check_field(field_name: str, annotation: object) -> Any:
    """Helper for validating a single schema-field.

    Recurses into the `annotation` to check that it is:
    - a normal type / typing annotation, or
    - a nested plain-dict schema mapping field names to annotations.

    Args:
        field_name (str): The name of the field to validate.
        annotation (object): The annotation to validate for the field.

    Returns:
        object: The same annotation if it's valid.

    Raises:
        InvalidConfigSchema: If:
            - the field name is not a non-empty string, or
            - the annotation is not a valid type hint, or
            - the annotation can't be resolved to a pydantic field type.
    """
    # field names must be non-empty strings
    if not isinstance(field_name, str) or not field_name:
        raise InvalidConfigSchema(f"Invalid field name {field_name!r}")

    # annotations can't be None or Ellipsis
    if annotation is None or annotation is ...:
        raise InvalidConfigSchema(
            f"Field {field_name!r} must map to a type/annotation, got {annotation!r}"
        )

    # distinguish between nested schemas & dicts containing values rather than valid
    # annotations
    if isinstance(annotation, dict):
        for nested_name, nested_ann in annotation.items():
            _check_field(nested_name, nested_ann)
        return annotation

    # Reject value-like runtime objects
    if not _is_annotation_like(annotation):
        raise InvalidConfigSchema(
            f"Field {field_name!r} must map to a type/annotation, got value-like "
            f"{annotation!r}"
        )

    try:  # sanity check: raises if annotation can't be resolved to pydantic field
        TypeAdapter(annotation)
    except Exception as e:
        raise InvalidConfigSchema(
            f"Invalid type hint for field {field_name!r}: {e}"
        ) from e

    return annotation


def _convert_dict(
    schema: PlainDictSchema,
    model_name: str,
) -> type[BaseModel]:
    """Construct BaseModel from a plain dictionary schema.

    Note:
        All parameters defined in a plain dictionary schema are considered required
        and extra fields are forbidden. This means that trying to instatiate a
        metric with missing or extra configuration parameters that are not defined
        in the schema will raise a validation error.

    Args:
        schema (PlainDictSchema): A plain dictionary mapping parameter names to
            their types.
        model_name (str): The name to give the dynamically created BaseModel.

    Returns:
        type[BaseModel]: The dynamically created BaseModel subclass.

    Raises:
        InvalidConfigSchema: If

            - the field name is not a non-empty string, or
            - the annotation is not a valid type hint, or
            - the annotation can't be resolved to a pydantic field type.
    """  # noqa: DOC502
    fields: dict[str, tuple[Any, Any]] = {}

    for field_name, raw_annotation in schema.items():
        # ensure valid field name & annotation
        annotation: Any = _check_field(field_name, raw_annotation)

        if isinstance(annotation, dict):  # convert nested dict to nested BaseModel
            nested_model: type[BaseModel] = _convert_dict(annotation, field_name)
            fields[field_name] = (nested_model, ...)
        else:  # convert normal annotation to pydantic field
            fields[field_name] = (annotation, ...)

    model = create_model(
        model_name, __config__=ConfigDict(extra="forbid"), **cast(Any, fields)
    )
    return cast(type[BaseModel], model)


def _unwrap_typeddict_annotation(annotation: object) -> Any:
    """Recursively remove Required[...] / NotRequired[...] markers from a TypedDict
    field annotation, preserving the underlying annotation structure.

    Args:
        annotation (object): The annotation to unwrap.

    Returns:
        Any: The unwrapped annotation with Required/NotRequired wrappers removed.
    """
    origin = get_origin(annotation)
    args = get_args(annotation)

    if origin in (Required, NotRequired):
        return _unwrap_typeddict_annotation(args[0])

    if not args:  # base case: not a wrapper, return as-is
        return annotation

    # recursively unwrap any args that are wrapped in Required/NotRequired
    stripped_args = tuple(_unwrap_typeddict_annotation(arg) for arg in args)

    if origin is None:  # pragma: no cover
        return annotation

    # reconstruct unions by reapplying the | operator to the unwrapped args
    if origin in (Union, UnionType):  # pragma: no cover
        result = stripped_args[0]
        for arg in stripped_args[1:]:
            result = result | arg
        return result

    # if the origin is a valid generic type, reconstruct it with the unwrapped args
    try:
        return (
            origin[stripped_args]
            if len(stripped_args) > 1
            else origin[stripped_args[0]]
        )
    except Exception:  # pragma: no cover
        # shouldn't happen; primarily defensive
        return annotation


def _convert_typeddict(td_cls: TypedDictSchema, model_name: str) -> type[BaseModel]:
    """Construct BaseModel from TypedDict.

    Note:
        The required/optional status of parameters is determined based on the
        `TypedDict`'s `__total__`, `__required_keys__`, and `__optional_keys__`
        attributes. Extra fields are forbidden, meaning that trying to instatiate a
        metric with missing or extra configuration parameters that are not defined
        in the schema will raise a validation error.

    Args:
        td_cls (TypedDictSchema): The TypedDict subclass to convert.
        model_name (str): The name to give the dynamically created BaseModel.

    Returns:
        type[BaseModel]: The dynamically created BaseModel subclass.

    Raises:
        InvalidConfigSchema: If

            - the field name is not a non-empty string, or
            - the annotation is not a valid type hint, or
            - the annotation can't be resolved to a pydantic field type.
    """  # noqa: DOC502
    # get names->types mapping and required/optional status
    hints: dict[str, Any] = get_type_hints(td_cls, include_extras=True)
    total: bool = bool(getattr(td_cls, "__total__", True))
    required_keys: set[str] = set(getattr(td_cls, "__required_keys__", set()))
    optional_keys: set[str] = set(getattr(td_cls, "__optional_keys__", set()))

    # construct the fields dict for create_model,
    # determining default values based on required/optional status
    fields: dict[str, tuple[Any, Any]] = {}
    for field_name, typeddict_annotation in hints.items():
        raw_annotation = _unwrap_typeddict_annotation(typeddict_annotation)
        annotation: Any = _check_field(field_name, raw_annotation)

        # determine if the field is required or optional
        if total:
            is_required = field_name not in optional_keys
        else:
            is_required = field_name in required_keys

        default = ... if is_required else None
        fields[field_name] = (annotation, default)

    model = create_model(
        model_name, __config__=ConfigDict(extra="forbid"), **cast(Any, fields)
    )
    return cast(type[BaseModel], model)


def _convert_to_pydantic_model(
    metric_cls: type[BaseMetric[Any]],
    schema: PlainDictSchema | TypedDictSchema,
) -> type[BaseModel]:
    """Helper for dynamically creating a pydantic model from a dictionary schema.

    Args:
        metric_cls (type[BaseMetric[Any]]): The metric class the schema is being
            applied to.
        schema (PlainDictSchema | TypedDictSchema): The schema to use for validation.

    Returns:
        pydantic_model (type[BaseModel]): The dynamically created pydantic model.

    Raises:
        InvalidConfigSchema: If

            - the field name is not a non-empty string, or
            - the annotation is not a valid type hint, or
            - the annotation can't be resolved to a pydantic field type.
        TypeError: If the provided `schema` argument is not a plain dictionary or a
            `TypedDict` subclass.
    """
    try:  # propagate validation errors as InvalidConfigSchema for better context
        if _is_typeddict_type(schema):
            return _convert_typeddict(schema, f"{metric_cls.__name__}ConfigModel")
        elif isinstance(schema, dict):
            return _convert_dict(schema, f"{metric_cls.__name__}ConfigModel")
        else:  # pragma: no cover
            raise TypeError(
                "Schema must be a plain dictionary or a TypedDict subclass."
            )
    except InvalidConfigSchema:
        raise


def _ensure_no_extras_and_frozen(config_schema: PydanticSchema) -> PydanticSchema:
    """Helper for ensuring that a pydantic model doesn't allow extra fields and is frozen.

    Args:
        config_schema (PydanticSchema): The pydantic schema to check and modify if
            necessary.

    Returns:
        PydanticSchema: The same pydantic schema with `extra="forbid"` and `frozen=True`
            in its Config.
    """
    config_schema.model_config["extra"] = "forbid"
    config_schema.model_config["frozen"] = True
    config_schema.model_rebuild(force=True)
    return config_schema


def config_schema(
    *, schema: PlainDictSchema | TypedDictSchema | PydanticSchema
) -> Callable[[type[BaseMetric[Any]]], type[BaseMetric[Any]]]:
    """Decorator for defining the expected schema of a metric's runtime parameters.

    The specified schema is converted to a pydantic model and set as the value of the
    decorated classes'
    [`config_schema`][evallm.metrics.base.BaseMetric.config_schema] class
    variable. This can subsequently be used for validating runtime parameters that are
    provided on instantiation. (see also:
    [`BaseMetric.config`][evallm.metrics.base.BaseMetric.config])

    ??? example annotate

        **Specifying a config schema**

        ...as a *plain dictionary*:
        ```python
        schema: PlainDictSchema = {"threshold": int}  # (1)!
        ```

        ...as a *`TypedDict` subclass*:
        ```python
        from typing import TypedDict, NotRequired


        class CustomConfig(TypedDict):
            # required parameter
            keywords: list[str]

            # optional parameter with default value of None
            min_keyword_counts: NotRequired[dict[str, int]]


        schema: TypedDictSchema = CustomConfig  # (2)!
        ```

        ...as a *`pydantic.BaseModel` subclass*:
        ```python
        from pydantic import BaseModel, Field, field_validator


        class ConfigWithValidation(BaseModel):
            # required parameter with validation
            keywords: list[str] = Field(..., min_length=1)

            # optional parameter with validation
            threshold: float | None = Field(default=None, ge=0, le=1)

            @field_validator("keywords")
            @classmethod
            def validate_keywords(cls, keywords: list[str]) -> list[str]:
                if not all(k.startswith("foo") for k in keywords):
                    raise ValueError("All keywords must start with 'foo'.")
                return keywords


        schema: PydanticSchema = ConfigWithValidation  # (3)!
        ```

        **Decorating a metric class with a config schema**

        ```python
        @register_metric
        @config_schema(schema=schema)  # (4)!
        class KeywordRatios(BaseMetric[dict[str, float]]):
            def evaluate(self, text: str) -> dict[str, float]:
                counts: dict[str, int] = {}
                for keyword in self.config.get("keywords"):  # (5)!
                    counts[keyword] = text.count(keyword)

                ratios: dict[str, float] = {}
                for keyword, count in counts.items():
                    kw_ratio: float = counts[keyword] / sum(counts.values())
                    if (threshold := self.config.get("threshold")) is not None:  # (6)!
                        if kw_ratio < threshold:
                            kw_ratio = 0.0
                    ratios[keyword] = kw_ratio

                return self.result(ratios)
        ```

    1.  A plain dictionary is the simplest way to define a schema, it will however
        assume all parameters as required.
    2.  A `TypedDict` subclass allows you to specify optional parameters using
        `typing.NotRequired` and provides better readability.
    3.  A [`pydantic.BaseModel`](https://docs.pydantic.dev/2.12/api/base_model/)
        subclass provides the most flexibility, allowing you to define complex
        validation logic and default values. The defined model will be used as-is.
    4.  The `@config_schema` decorator should be applied *before* the
        `@register_metric` decorator.
    5.  You can access validated configuration *values* in the metric's `evaluate(...)`
        method via `self.config.get(<parameter_name>)`.
    6.  Remember to appropriately handle parameters that you defined as optional in the
        schema!

    Args:
        schema (PlainDictSchema | TypedDictSchema | PydanticSchema): The schema to
            enforce for the decorated metric class. Must be one of the following:

            - A plain dictionary mapping parameter names to their types
            - A `TypedDict` subclass
            - A `pydantic.BaseModel` subclass

    Returns:
        metric_cls (Callable[[type[BaseMetric[Any]]], type[BaseMetric[Any]]]): The decorated metric class.
    """

    def decorate(
        metric_cls: type[BaseMetric[Any]],
    ) -> type[BaseMetric[Any]]:
        """Decorate a metric class by setting its `_config_schema` class variable to
        the provided schema.

        Args:
            metric_cls (type[BaseMetric[Any]]): The metric class to
                decorate.

        Returns:
            type[BaseMetric[Any]]: The decorated metric class with
                its `_config_schema` attribute set to the provided schema as a
                pydantic model.

        Raises:
            TypeError: If the provided `schema` argument is not one of:

                - a plain dictionary, mapping parameter names to their types, or
                - a `TypedDict` subclass, or
                - a `pydantic.BaseModel` subclass.
            InvalidConfigSchema: If

                - the field name is not a non-empty string, or
                - the annotation is not a valid type hint, or
                - the annotation can't be resolved to a pydantic field type.
        """
        if isinstance(schema, type) and issubclass(schema, BaseModel):
            schema_model: PydanticSchema = _ensure_no_extras_and_frozen(schema)
            metric_cls.config_schema = schema_model
            return metric_cls
        elif _is_typeddict_type(schema) or isinstance(schema, dict):  # pragma: no cover
            try:  # convert plain dict or TypedDict schema to a pydantic model
                schema_model = _convert_to_pydantic_model(metric_cls, schema)
                metric_cls.config_schema = _ensure_no_extras_and_frozen(schema_model)
                return metric_cls
            except InvalidConfigSchema:
                raise
        else:  # pragma: no cover
            raise TypeError(
                "Schema must be one of: a plain dictionary, a TypedDict subclass, or "
                + "a pydantic BaseModel subclass."
            )

    return decorate


__all__ = [
    "InvalidConfigSchema",
    "PlainDictSchema",
    "PydanticSchema",
    "TypedDictSchema",
    "config_schema",
]
