#!/usr/bin/env python3

"""![status](../../../badges/status-application-models.svg) Canonical internal data contracts as DTOs

DTOs are implemented as immutable
[:simple-pydantic: `BaseModel`](https://docs.pydantic.dev/2.12/api/base_model/)s that
perform *structural validation only*. While this allows them to stay decoupled from
other layers, it also means that *no semantic/ stateful validation* is performed.
"""

from __future__ import annotations

from abc import ABC
from collections.abc import Mapping
from pathlib import Path
from typing import Annotated, Any, ClassVar, Literal, Self, cast

from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    FilePath,
    JsonValue,
    ModelWrapValidatorHandler,
    StringConstraints,
    field_validator,
    model_validator,
)

from evallm.application.errors import (
    MetricConfigurationError,
    UnknownMetricError,
)
from evallm.metrics.base import BaseMetric
from evallm.metrics.registry import Catalog
from evallm.utils import make_crossref_resolver


class BaseDTO(BaseModel, ABC):
    """Abstract base class for all DTOs in the application layer."""

    model_config = ConfigDict(
        # disallow extra fields
        extra="forbid",
        # make all DTOs immutable by default
        frozen=True,
        # resolve mkdocstrings-style cross-references in JSON schema descriptions to URLs
        json_schema_extra=make_crossref_resolver(
            docs_base_url="https://psandhaas.github.io/evaLLM/reference"
        ),
        # include fields with default values in 'required' during JSON-serialization
        # (this does not affect model_json_schema(mode='validation'))
        json_schema_serialization_defaults_required=True,
        # ensure validation for all DTOs
        revalidate_instances="subclass-instances",
        # keep NaN and Infinity as-is in JSON output, instead of converting to null
        ser_json_inf_nan="constants",
        # use source code docstrings for field descriptions in JSON schema
        use_attribute_docstrings=True,
    )

    def __new__(cls, *args, **kwargs) -> Self:  # noqa: ARG004
        if cls is BaseDTO:
            raise NotImplementedError(
                f"{cls.__name__} is an abstract base class and cannot be instantiated "
                + "directly."
            )
        return super().__new__(cls)

    @classmethod
    def create(
        cls,
        obj: Any | None = None,
        *,
        json_payload: str | bytes | bytearray | None = None,
        **kwargs: Any,
    ) -> Self:
        """Convenience factory to create DTOs from various input sources.

        Args:
            obj (Any | None): The DTO's fields as a native Python object (e.g. a dict).
            json_payload (str | bytes | bytearray | None): The DTO's fields as a JSON
                string- or bytes-payload.
            **kwargs: Keyword arguments corresponding to the DTO's fields.

        Returns:
            Self: An instance of the DTO created from the provided input.

        Raises:
            ValueError: If 0 or more than 1 of the input sources are provided.
        """
        if (
            sum(
                [
                    obj is not None,
                    json_payload is not None,
                    bool(kwargs),
                ]
            )
            != 1
        ):
            raise ValueError(
                "Provide exactly one of: obj, json_payload, or keyword arguments."
            )

        if json_payload is not None:
            return cls.model_validate_json(json_payload)

        if obj is not None:
            return cls.model_validate(obj)

        return cls.model_validate(kwargs)


class MetricSpec(BaseDTO):
    """DTO specifying a metric and its configuration"""

    name: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
    """Non-empty name under which a metric is registered.

    !!! note "See also"
        - [`BaseMetric.name`][evallm.metrics.base.BaseMetric.name]
        - [`@register_metric`][evallm.metrics.registry.register_metric]
        - [`MetricsCatalog.available`][evallm.metrics.registry.MetricsCatalog.available]
    """
    config: dict[str, JsonValue] | None = None
    """Runtime configuration for the metric or `None` (default), if not needed.

    !!! note "See also"
        - [`BaseMetric.config_schema`][evallm.metrics.base.BaseMetric.config_schema]
        - [`BaseMetric.from_config`][evallm.metrics.base.BaseMetric.from_config]
    """
    result_key: (
        Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] | None
    ) = None
    """Optional key to use in the evaluation result dictionary for this metric
    instance's value."""

    def resolve_metric(self, catalog: type[Catalog]) -> BaseMetric:
        """Resolve the metric specification to a
        [`BaseMetric`][evallm.metrics.base.BaseMetric] instance.

        Performs contextual validation of the DTO by:

        1. looking up the specified [`name`][..name] in the provided
        [`catalog`][evallm.metrics.registry.Catalog] and
        2. trying to instantiate the metric with the provided [`config`][..config].

        If either of these steps fail, an appropriate error is raised.

        !!! note "Metric instance's `result_key`"
            The resolved metric instance's `result_key` attribute is **not** yet set to
            the value of this DTO's `result_key` field. This is because at this point
            there's no guarantee that the provided `result_key` is unique among all
            resolved metrics for a given evaluation request.

        Args:
            catalog (type[Catalog]): A catalog class to use for looking up the metric.

        Returns:
            BaseMetric: An instance of the specified metric, if resolution and
                validation succeed.

        Raises:
            TypeError: If the provided catalog does not implement the
                [`Catalog`][evallm.metrics.registry.Catalog] interface.
            UnknownMetricError: If the specified metric name is not registered in the
                catalog.
            MetricConfigurationError: If the metric cannot be instantiated from the
                provided config.
        """
        if not issubclass(catalog, Catalog):
            raise TypeError(
                f"Provided catalog {catalog.__name__} doesn't implement the Catalog "
                + "interface"
            )
        if self.name not in (available := catalog.available()):
            raise UnknownMetricError(metric_name=self.name, available=available)
        metric_cls: type[BaseMetric] = catalog.get_metric(self.name)
        try:
            return metric_cls.from_config(**cast(dict[str, Any], self.config or {}))
        except ValueError as exc:
            raise MetricConfigurationError(
                metric_name=self.name,
                config=self.config,
                reason=str(exc),
            ) from exc


class BaseDatasetSpec(BaseDTO, ABC):
    """Abstract base class for dataset specifications."""

    _specs_by_format: ClassVar[dict[str, type[BaseDatasetSpec]]] = {}
    """Registry of known dataset spec types by discriminator value."""

    data_format: Annotated[
        str,
        StringConstraints(strip_whitespace=True, min_length=1, to_lower=True),
    ]
    """Dataset format discriminator used for dataset spec resolution."""

    def __new__(cls, *args, **kwargs) -> Self:  # noqa: ARG004
        if cls is BaseDatasetSpec:
            raise NotImplementedError(
                f"{cls.__name__} is an abstract base class and cannot be instantiated "
                + "directly."
            )
        return super().__new__(cls)

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        if cls is BaseDatasetSpec:
            return

        data_format = getattr(cls, "data_format", None)
        if isinstance(data_format, str) and data_format:
            BaseDatasetSpec._specs_by_format[data_format] = cls

    @classmethod
    def infer_data_format(cls, data: Mapping[str, Any]) -> str | None:
        """Infer dataset format from the input-file suffix, if possible.

        Args:
            data (Mapping[str, Any]): Raw dataset payload.

        Returns:
            str | None: Inferred dataset format, or `None` when no inference is
                possible.
        """
        path_like = data.get("input_file")
        if path_like is None:  # pragma: no cover
            return None
        suffix = Path(str(path_like)).suffix.strip().lower().lstrip(".")
        if not suffix:
            return None
        return suffix

    @classmethod
    def resolve(
        cls,
        data: BaseDatasetSpec | Mapping[str, Any],
    ) -> BaseDatasetSpec:
        """Resolve a raw or already-validated dataset payload to a concrete dataset DTO.

        Args:
            data (BaseDatasetSpec | Mapping[str, Any]): Raw or typed dataset payload.

        Returns:
            BaseDatasetSpec: Resolved dataset spec instance.

        Raises:
            ValueError: If the dataset format is missing, unsupported, or conflicts
                with the inferred format.
        """
        if isinstance(data, BaseDatasetSpec):  # pragma: no cover
            return data

        payload: dict[str, Any] = dict(data)
        explicit_value = payload.get("data_format")
        explicit: str | None = None
        if explicit_value is not None:
            if (
                not isinstance(explicit_value, str) or not explicit_value.strip()
            ):  # pragma: no cover
                raise ValueError("dataset.data_format must be a non-empty string")
            explicit = explicit_value.strip().lower()

        inferred = cls.infer_data_format(payload)
        if explicit is not None and inferred is not None and explicit != inferred:
            raise ValueError(
                "dataset.data_format does not match inferred format from input_file "
                + f"(explicit={explicit!r}, inferred={inferred!r})"
            )

        resolved_format = explicit or inferred
        if resolved_format is None:
            raise ValueError(
                "Could not infer dataset.data_format from input_file. "
                + "Provide dataset.data_format explicitly."
            )

        if resolved_format not in cls._specs_by_format:
            available = ", ".join(sorted(cls._specs_by_format))
            raise ValueError(
                f"Unsupported dataset format {resolved_format!r}. "
                + f"Available formats: {available}."
            )

        payload["data_format"] = resolved_format
        dataset_cls = cls._specs_by_format[resolved_format]
        return dataset_cls.model_validate(payload)


class JsonlDatasetSpec(BaseDatasetSpec):
    """DTO specifying a structured JSONL dataset to evaluate"""

    data_format: Literal["jsonl"] = "jsonl"
    """Dataset format discriminator."""

    input_file: FilePath
    """Path to a local JSONL-file containing the dataset to evaluate."""
    record_id_key: Annotated[
        str, StringConstraints(strip_whitespace=True, min_length=1)
    ] = "#IDX#"
    """The dataset field/ column name that contains the unique identifier for each
    record, specified as a non-empty string. Defaults to `"#IDX#"`, which indicates
    that record IDs should be generated based on the record's (0-indexed) position in
    the dataset."""
    text_key: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
    """The dataset field/ column name that contains the text to evaluate, specified as
    a non-empty string."""

    @field_validator("input_file")
    @classmethod
    def _validate_jsonl_path(cls, path: FilePath) -> FilePath:
        """Validator to check if the input file has a .jsonl extension

        Args:
            path (FilePath): The path to validate

        Returns:
            FilePath: The validated path if it has a .jsonl extension

        Raises:
            ValueError: If the input file does not have a .jsonl extension
        """
        if not path.name.endswith(".jsonl"):
            raise ValueError("Input file must have a .jsonl extension")
        return path

    @model_validator(mode="wrap")
    @classmethod
    def _validate_keys(cls, data: Any, handler: ModelWrapValidatorHandler) -> Self:
        """Model-level validator to ensure that `record_id_key` and `text_key` are not
        the same, as this would cause a conflict when trying to extract the record ID
        and text from each dataset record.

        Args:
            data (Any): The input data to validate.
            handler (ModelWrapValidatorHandler): The pydantic handler for wrap validators.

        Returns:
            Self: The validated model instance if validation passes.

        Raises:
            ValueError: If `text_key` is missing or if it is the same as
                `record_id_key`.
        """
        if isinstance(data, str):  # pragma: no cover
            data = cls.model_validate_json(data)
        if isinstance(data, BaseModel):  # pragma: no cover
            data = data.model_dump()
        if (text_key := data.get("text_key")) is None:
            raise ValueError("text_key is required")
        if data.get("record_id_key", "#IDX#") == text_key:
            raise ValueError("record_id_key and text_key cannot be the same")
        return cast(Self, handler(data))


class EvaluationRequest(BaseDTO):
    """DTO specifying which configuration of metrics to evaluate a given dataset on"""

    dataset: BaseDatasetSpec
    """The dataset to evaluate, specified as a concrete
    [`BaseDatasetSpec`][evallm.application.models.BaseDatasetSpec] subtype."""
    metrics: Annotated[list[MetricSpec], Field(min_length=1)]
    """List of one or more metrics to use for evaluation, each specified as a
    [`MetricSpec`][evallm.application.models.MetricSpec]."""

    @field_validator("dataset", mode="plain")
    @classmethod
    def _validate_dataset_spec(
        cls,
        value: Any,
    ) -> BaseDatasetSpec:
        """Resolve dataset payloads to concrete dataset spec DTOs.

        Behavior:
            - use explicit ``data_format`` when provided
            - infer ``data_format`` from ``input_file`` suffix when omitted
            - raise when explicit and inferred formats conflict
            - raise when no format can be inferred and none is provided

        Returns:
            BaseDatasetSpec: A resolved concrete dataset spec.

        Raises:
            TypeError: If the input value is neither a dataset spec nor a mapping.
        """
        if isinstance(value, BaseDatasetSpec):
            return value
        if isinstance(value, Mapping):
            return BaseDatasetSpec.resolve(value)
        raise TypeError("dataset must be a dataset spec object or mapping")


class DatasetRecord(BaseDTO):
    """DTO representing a single record to evaluate in a dataset"""

    record_id: str
    """The record's unique identifier, specified as a string. The value of this field
    corresponds to the dataset field specified by
    [`JsonlDatasetSpec.record_id_key`][evallm.application.models.JsonlDatasetSpec.record_id_key]."""
    text: str
    """The record's value to evaluate. The value of this field corresponds to the
    dataset field specified by
    [`JsonlDatasetSpec.text_key`][evallm.application.models.JsonlDatasetSpec.text_key],
    and indexed by [`record_id`][..record_id].

    !!! note "Type coercion for text field"
        During validation, `None` values will be converted to the empty string
        (i.e. `""`), values of JSON-serializable types that aren't strings will be
        converted to their string representation (e.g. `{'key': 123}` will become
        `"{'key': 123}"`), and string values are kept as-is (e.g. `" "` will remain
        `" "`).
    """
    metadata: dict[str, JsonValue] | None = None
    """Optional dictionary containing any additional fields from the dataset record."""

    @field_validator("text", mode="before", json_schema_input_type=JsonValue)
    @classmethod
    def _stringify_text(cls, v: JsonValue) -> str:
        """Convert the record's value to a string or raise a validation error if this
        isn't possible.

        Args:
            v (JsonValue): The value as any JSON-serializable type.

        Returns:
            str: The value converted to a string.

        Raises:
            ValueError: If the value cannot be converted to a string.
        """
        if v is None:
            return ""
        if isinstance(v, str):
            return v
        try:
            return str(v)
        except Exception:
            raise ValueError(
                f"Value {v!r} of type {type(v).__name__} cannot be converted to string"
            ) from None


class EvaluationResult(BaseDTO):
    """DTO representing the evaluated metrics' results for a single record"""

    record_id: str
    """The unique identifier for the evaluated record, specified as a string.
    The value of this field corresponds to the dataset field specified by
    [`JsonlDatasetSpec.record_id_key`][evallm.application.models.JsonlDatasetSpec.record_id_key]
    and mirrors the value of the corresponding
    [`DatasetRecord.record_id`][evallm.application.models.DatasetRecord.record_id]."""
    results: dict[str, JsonValue]
    """A dictionary mapping metric result keys to their computed results.

    Each key corresponds to the resolved metric instance's
    [`result_key`][evallm.metrics.base.BaseMetric.result_key], which defaults to the
    metric's [`name`][evallm.metrics.base.BaseMetric.name] unless explicitly
    overridden.

    !!! note "See also"
        - [`MetricSpec.result_key`][evallm.application.models.MetricSpec.result_key]
        - [`BaseMetric.result_key`][evallm.metrics.base.BaseMetric.result_key]
        - [`CompositeMetric.evaluate`][evallm.metrics.composite.CompositeMetric.evaluate]
    """
