#!/usr/bin/env python3

"""`JsonFormatMetric` — validate the adherence of a text to an expected output format

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

from __future__ import annotations

__all__ = []

from typing import Any

import orjson
from jsonschema import Draft202012Validator, SchemaError
from pydantic import BaseModel, Field, field_validator

from evallm.metrics.base import BaseMetric
from evallm.metrics.config import config_schema
from evallm.metrics.registry import register_metric


class JsonFormatConfig(BaseModel):
    """The configuration schema for the [`JsonFormatMetric`][evallm.metrics.builtin.json_format.JsonFormatMetric].

    Keyword Args:
        expected_schema (str | dict[str, Any]): A valid JSON schema conforming to the
            [JSON Schema draft-2020-12 specification](https://json-schema.org/draft/2020-12/json-schema-core)
            that describes the expected output format. Can be provided either as a
            JSON payload string or as a dictionary. Note that the schema must be a
            top-level schema, defined as `'type': 'object'` with a `'properties'`
            attribute containing the expected fields. Fields that aren't explicitly
            listed in the `'required'` attribute of the schema, are *considered
            optional by default*.
        check_formats (bool, optional): (default=`False`) Whether to check constraints
            in 'format' annotations of primitive types. If `True`, the following
            formats will additionally be validated where present

            - `'date'`
            - `'date-time'`
            - `'duration'`
            - `'email'`
            - `'hostname'`
            - `'idn-email'`
            - `'idn-hostname'`
            - `'ipv4'`
            - `'ipv6'`
            - `'iri'`
            - `'iri-reference'`
            - `'json-pointer'`
            - `'regex'`[^1]
            - `'relative-json-pointer'`
            - `'time'`
            - `'uri'`
            - `'uri-reference'`
            - `'uri-template'`
            - `'uuid'`

            [^1]:
                While the JSON Schema specification
                [recommends](https://json-schema.org/draft/2020-12/json-schema-core#name-regular-expressions)
                using ECMA 262 regular expressions, validation for format annotations like
                `'pattern'` or `'patternProperties'` will use Python regular expressions
                instead.

    ??? example "Example configuration"
        ```python
        json_payload: str = '''
        {
            "type": "object",
            "properties": {
                "question": {
                    "type": "string"
                },
                "setup": {
                    "type": "string"
                },
                "joke": {
                    "type": "string"
                },
                "funniness_score": {
                    "type": "number",
                    "minimum": 0.0,
                    "maximum": 1.0
                }
            },
            "required": ["question", "joke"]
        }
        '''
        config: dict[str, Any] = {"expected_schema": json_payload}
        metric: JsonFormatMetric = JsonFormatMetric.from_config(**config)
        ```
    """

    expected_schema: dict[str, Any] = Field(
        ...,
        description=str(
            "A valid JSON schema that conforms to the JSON Schema draft-2020-12 "
            + "specification (see "
            + "https://json-schema.org/draft/2020-12/json-schema-core) as a "
            + "dictionary, describing the expected structured output format."
        ),
    )
    check_formats: bool = Field(
        default=False,
        description=str(
            "Whether constraints in 'format' annotations should be checked. If "
            + "`True`, the following formats will additionally be checked where "
            + "present: `'date'`, `'date-time'`, `'duration'`, `'email'`, "
            + "`'hostname'`, `'idn-email'`, `'idn-hostname'`, `'ipv4'`, `'ipv6'`, "
            + "`'iri'`, `'iri-reference'`, `'json-pointer'`, `'regex'`, "
            + "`'relative-json-pointer'`, `'time'`, `'uri'`, `'uri-reference'`, "
            + "`'uri-template'`, `'uuid'`."
        ),
    )

    @field_validator("expected_schema", mode="before")
    @classmethod
    def _validate_expected_schema(cls, data: str | dict[str, Any]) -> dict[str, Any]:
        """Validate that `expected_schema` is a valid schema conforming to the [JSON Schema draft-2020-12 specification](https://json-schema.org/draft/2020-12/json-schema-core).

        Args:
            data (str | dict[str, Any]): The JSON Schema to validate, either as a
                JSON string or a dictionary.

        Returns:
            expected_schema (dict[str, Any]): The validated JSON Schema as a dictionary

        Raises:
            ValueError: If
                - the provided schema is not a valid JSON Schema according to
                    the draft-2020-12 specification,
                - the provided schema is not a top-level object schema, or
                - the provided schema does not define a `properties` field.
        """
        if isinstance(data, str):
            try:
                schema: dict[str, Any] = orjson.loads(data)
            except orjson.JSONDecodeError as e:
                raise ValueError(
                    "Invalid JSON-string. Failed to parse expected_schema."
                ) from e
        else:
            schema = data
        try:
            Draft202012Validator.check_schema(schema)
        except SchemaError as e:
            raise ValueError(
                "Invalid schema. The provided schema must conform to the draft-2020-12 "
                + "JSON Schema specification."
            ) from e

        if schema.get("type") != "object":
            raise ValueError("expected_schema must be a top-level object schema")

        if "properties" not in schema or not isinstance(schema["properties"], dict):
            raise ValueError("expected_schema must define 'properties'")

        return schema


@register_metric
@config_schema(schema=JsonFormatConfig)
class JsonFormatMetric(BaseMetric[dict[str, bool]]):  # type: ignore[type-var]
    """Metric for evaluating whether text matches a given JSON format.

    This metric evaluates whether the input text adheres to a structured output format,
    defined by the `expected_schema`.

    Keyword Args:
        expected_schema (dict[str, Any]): The expected JSON schema as a dictionary
            that conforms to the [JSON Schema draft-2020-12 specification](https://json-schema.org/draft/2020-12/json-schema-core).
        check_formats (bool, optional): If `True`, format annotations are also
            validated. Defaults to `False`.
    """

    def evaluate(self, text: str) -> dict[str, dict[str, bool]]:
        """Evaluate whether the input text matches the expected JSON format.

        This method attempts to parse the input text as JSON and checks whether it
        contains the expected fields as specified in the configuration.

        Args:
            text (str): The input text to evaluate.

        Returns:
            dict[str, dict[str, bool]]: A dictionary mapping this metric instance's
                [`result_key`][evallm.metrics.base.BaseMetric.result_key] to the
                evaluation results, where each expected field is mapped to a boolean,
                indicating whether that field is valid. Note that all fields will be
                `False` if the input text is not valid JSON.
        """
        expected_schema: dict[str, Any] = self.config["expected_schema"]
        res: dict[str, bool] = {field: False for field in expected_schema["properties"]}

        # handle empty inputs
        if text is None or text.strip() == "":
            return self.result(res)

        try:  # parse input text as JSON
            parsed: dict[str, Any] = orjson.loads(text)
        except orjson.JSONDecodeError:
            return self.result(res)

        # validate fields according to the expected schema
        for field, field_schema in expected_schema["properties"].items():
            if self.config["check_formats"]:
                validator: Draft202012Validator = Draft202012Validator(
                    schema=field_schema,
                    format_checker=Draft202012Validator.FORMAT_CHECKER,
                )
            else:
                validator = Draft202012Validator(schema=field_schema)

            field_value: Any = parsed.get(field)
            required: bool = field in expected_schema.get("required", [])
            if field in parsed:
                if not required:
                    if field_value is None:  # optional field with null value -> valid
                        res[field] = True
                    else:  # optional field is present -> check validity
                        res[field] = validator.is_valid(parsed[field])
                else:  # check validity of required field
                    res[field] = validator.is_valid(parsed[field])
            elif not required:  # missing optional fields are valid
                res[field] = True

        return self.result(res)
