#!/usr/bin/env python3

"""![status](../../../badges/status-application-services.svg) Application services for evaLLM

This module contains services of the evaLLM application layer, responsible for
coordinating other components.
"""

from __future__ import annotations

from collections.abc import Iterator
from contextlib import AbstractContextManager
from typing import Any, cast

from pydantic import JsonValue

from evallm.application.adapters import resolved_records
from evallm.application.errors import (
    DatasetAccessError,
    DuplicateMetricResultKeyError,
    MetricConfigurationError,
    UnknownMetricError,
    UnsupportedDatasetError,
)
from evallm.application.models import (
    BaseDatasetSpec,
    DatasetRecord,
    EvaluationRequest,
    EvaluationResult,
    JsonlDatasetSpec,
)
from evallm.application.validation import ensure_unique_metric_result_keys
from evallm.metrics.base import BaseMetric
from evallm.metrics.composite import CompositeMetric
from evallm.metrics.registry import Catalog, MetricsCatalog
from evallm.readers.reader import BaseReader, JsonlReader, ReaderError


class RequestResolver:
    """Component responsible for resolving & validating evaluation requests"""

    def __init__(self, *, catalog: type[Catalog] = MetricsCatalog) -> None:
        assert issubclass(
            catalog, Catalog
        ), "Provided catalog doesn't implement the Catalog interface"
        self._catalog = catalog

    def __call__(
        self, request: EvaluationRequest
    ) -> tuple[AbstractContextManager[Iterator[DatasetRecord]], CompositeMetric]:
        """Validate & resolve the provided evaluation request.

        This method returns an iterator over resolved
        [`DatasetRecord`][evallm.application.models.DatasetRecord]s and a
        [`CompositeMetric`][evallm.metrics.composite.CompositeMetric] instance
        containing the resolved metrics to use for evaluation.

        Args:
            request (EvaluationRequest): The evaluation request to validate.

        Returns:
            tuple[AbstractContextManager[Iterator[DatasetRecord]], CompositeMetric]: A
                tuple containing the resolved dataset records and the composite metric.

        Raises:
            DatasetAccessError: If the dataset cannot be accessed or read.
            UnsupportedDatasetError: If no reader can handle the specified dataset.
            UnknownMetricError: If a specified metric can't be found in the catalog.
            MetricConfigurationError: If a metric can't be instantiated from the
                provided runtime config.
            DuplicateMetricResultKeyError: If multiple metrics in the request would
                resolve to the same final result key.
        """
        dataset: AbstractContextManager[Iterator[DatasetRecord]]
        metrics: CompositeMetric
        try:
            dataset = self._resolved_dataset(request)
            metrics = self._compose_metrics(request)
            return dataset, metrics
        except (
            DatasetAccessError,
            UnsupportedDatasetError,
            DuplicateMetricResultKeyError,
            UnknownMetricError,
            MetricConfigurationError,
        ):
            raise

    def _create_reader(self, spec: BaseDatasetSpec) -> BaseReader[dict[str, Any]]:
        """Helper for creating a reader, capable of reading the specified dataset.

        ??? note "Implementational note"
            For now, we only support JSONL datasets, but this method can be extended in
            the future to support additional dataset types by checking the type of
            `spec` and returning an appropriate reader instance.

        Args:
            spec (BaseDatasetSpec): The dataset specification to create a reader for.

        Returns:
            BaseReader[dict[str, Any]]: A reader instance yielding raw dataset records
                as dictionaries.

        Raises:
            UnsupportedDatasetError: If no reader can handle the specified dataset.
        """
        if isinstance(spec, JsonlDatasetSpec):
            return JsonlReader(spec.input_file)
        raise UnsupportedDatasetError(type(spec))

    def _resolved_dataset(
        self, request: EvaluationRequest
    ) -> AbstractContextManager[Iterator[DatasetRecord]]:
        """Resolve the dataset specification in the request to an iterator of
        [`DatasetRecord`][evallm.application.models.DatasetRecord] instances.

        Args:
            request (EvaluationRequest): The evaluation request containing the dataset
                specification to resolve.

        Returns:
            AbstractContextManager[Iterator[DatasetRecord]]: A context manager yielding an
                iterator of resolved dataset records, each uniquely identified through its
                [`record_id`][evallm.application.models.DatasetRecord.record_id]
                and containing a [`text`][evallm.application.models.DatasetRecord.text]
                field with the content to evaluate.

        Raises:
            UnsupportedDatasetError: If no reader can handle the specified dataset.
            DatasetAccessError: If the dataset cannot be accessed or read.
        """
        try:
            if not isinstance(request.dataset, JsonlDatasetSpec):
                raise UnsupportedDatasetError(type(request.dataset))
            reader: BaseReader[dict[str, Any]] = self._create_reader(request.dataset)
            return resolved_records(reader, spec=request.dataset)
        except (FileNotFoundError, ReaderError) as exc:
            raise DatasetAccessError(
                path=str(getattr(request.dataset, "input_file", "<unknown>")),
                reason=str(exc),
            ) from exc
        except UnsupportedDatasetError:
            raise

    def _ensure_unique_result_keys(self, request: EvaluationRequest) -> None:
        """Ensure that all metric instances resolved from the request have unique
        result keys, raising an error if not.

        This is necessary because the evaluation results for each record are returned
        as a single dictionary mapping metrics'
        [`result_key`][evallm.metrics.base.BaseMetric.result_key] to their computed
        values. If multiple metrics resolve to the same result key, their values would
        overwrite each other in this dictionary, making it impossible to distinguish
        them in the results.

        Args:
            request (EvaluationRequest): The evaluation request containing the metric
                specifications to check.

        Raises:
            DuplicateMetricResultKeyError: If multiple metrics are specified with the
                same result key in the request. This includes duplicate explicit
                `result_key` values as well as collisions caused by metrics falling
                back to their names when `result_key` is not provided.
        """  # noqa: DOC502
        ensure_unique_metric_result_keys(request.metrics)

    def _compose_metrics(self, request: EvaluationRequest) -> CompositeMetric:
        """Resolve the metric specifications in the request to
        [`BaseMetric`][evallm.metrics.base.BaseMetric] instances and combine them into
        a [`CompositeMetric`][evallm.metrics.composite.CompositeMetric].

        Before composing the metrics, this method ensures that all specified metrics
        would resolve to unique final result keys. It then resolves each metric,
        assigns the instance's [`result_key`][evallm.metrics.base.BaseMetric.result_key],
        and adds it to the composite.

        Args:
            request (EvaluationRequest): The evaluation request containing the metric
                specifications to resolve.

        Returns:
            CompositeMetric: A composite metric containing the resolved metric
                instances.

        Raises:
            UnknownMetricError: If a specified metric can't be found in the catalog.
            MetricConfigurationError: If a metric can't be instantiated from the provided
                runtime config.
            DuplicateMetricResultKeyError: If multiple metrics in the request would
                resolve to the same final result key.
        """
        try:  # ensure that no duplicate result keys are present
            self._ensure_unique_result_keys(request)
        except DuplicateMetricResultKeyError:
            raise

        composite: CompositeMetric = CompositeMetric()
        for spec in request.metrics:
            metric: BaseMetric
            try:
                metric = spec.resolve_metric(self._catalog)
                # assign the instance's result key
                metric.result_key = spec.result_key  # uses metric's name if None
                # add the resolved metric instance to the composite
                composite.add(metric)
            except (UnknownMetricError, MetricConfigurationError):
                raise
        return composite


class Evaluator:
    """Component responsible for orchestrating the evaluation process.

    Takes an [`EvaluationRequest`][evallm.application.models.EvaluationRequest] as
    initialization parameter, resolves it, and provides a [`run`][.run]
    method that executes the evaluation, lazily yielding individual
    [`EvaluationResult`][evallm.application.models.EvaluationResult]s for each record
    in the specified dataset.

    Args:
        request (EvaluationRequest): The evaluation request specifying the dataset and
            metrics to use for evaluation.

    Keyword Args:
        catalog (type[Catalog] | None): An optional catalog class to use for
            resolving metric specifications. Defaults to `MetricsCatalog`.
            (Keyword-only)

    ???+ example annotate "Usage"
        ```python
        from pathlib import Path
        from pprint import pprint

        from evallm.application import Evaluator
        from evallm.application.models import (
            EvaluationRequest,
            JsonlDatasetSpec,
            MetricSpec,
        )
        from evallm.metrics import builtin  # (1)!

        request: EvaluationRequest = EvaluationRequest.create(  # (2)!
            obj={  # (7)!
                "dataset": {  # (3)!
                    "input_file": "C:/Users/.../Opus-4.6-Reasoning-3000x-filtered.jsonl",
                    "record_id_key": "id",
                    "text_key": "thinking",
                },
                "metrics": [  # (4)!
                    {
                        "name": "CountsMetric",
                        "config": {
                            "segments": "tokens",
                        },
                    },
                    {
                        "name": "JsonFormatMetric",
                        "config": {
                            "expected_schema": {
                                "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"],
                            }
                        },
                    },
                ],
            }
        )

        evaluator: Evaluator = Evaluator(request)  # (5)!

        for i, res in enumerate(evaluator.run()):  # (6)!
            if i > 3:
                break
            pprint(res.model_dump(), sort_dicts=False)
        ```
        ??? example "Output"
            ```pycon
            {
                "record_id": "orca_3452",
                "results": {
                    "CountsMetric": 257,
                    "JsonFormatMetric": {
                        "question": False,
                        "setup": False,
                        "joke": False,
                        "funniness_score": False,
                    },
                },
            }
            {
                "record_id": "drive_minimax_m2.1_questions_220036",
                "results": {
                    "CountsMetric": 121,
                    "JsonFormatMetric": {
                        "question": False,
                        "setup": False,
                        "joke": False,
                        "funniness_score": False,
                    },
                },
            }
            {
                "record_id": "orc_73480",
                "results": {
                    "CountsMetric": 243,
                    "JsonFormatMetric": {
                        "question": False,
                        "setup": False,
                        "joke": False,
                        "funniness_score": False,
                    },
                },
            }
            {
                "record_id": "drive_minimax_m2.1_questions_184095",
                "results": {
                    "CountsMetric": 181,
                    "JsonFormatMetric": {
                        "question": False,
                        "setup": False,
                        "joke": False,
                        "funniness_score": False,
                    },
                },
            }
            ```

    1.  **Important**: Remember to import the module where the metrics that you're going to
        use are defined. Otherwise evaLLM won't be able to discover them!
    2.  Construct an [`EvaluationRequest`][evallm.application.models.EvaluationRequest]
        programatically.
    3.  Define the [dataset][evallm.application.models.EvaluationRequest.dataset] to
        evaluate. Here we specify a JSONL file, along with its
        [record ID][evallm.application.models.JsonlDatasetSpec.record_id_key] and
        [text][evallm.application.models.JsonlDatasetSpec.text_key] keys.
    4.  Define a list of metrics to compute. Each
        [`MetricSpec`][evallm.application.models.MetricSpec] needs a
        [name][evallm.application.models.MetricSpec.name] under which the metric is
        registered and - depending on the metric - can optionally include
        [runtime configuration][evallm.application.models.MetricSpec.config].
    5.  Create an [`Evaluator`][evallm.application.services.Evaluator] with the evaluation
        request.
    6.  Lazily iterate over the
        [`EvaluationResult`][evallm.application.models.EvaluationResult]s.
    7.  You could also post the same request as a JSON payload:
        ```json
        {
        "dataset": {
            "input_file": "C:/Users/.../Opus-4.6-Reasoning-3000x-filtered.jsonl",
            "record_id_key": "id",
            "text_key": "thinking"
        },
        "metrics": [
            {
            "name": "CountsMetric",
            "config": {
                "segments": "tokens"
            }
            },
            {
            "name": "JsonFormatMetric",
            "config": {
                "expected_schema": {
                "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"]
                }
            }
            }
        ]
        }
        ```

    ??? info "`EvaluationRequest` JSON schema"
        ```json
        {'$defs': { 'JsonValue': {},
                    'JsonlDatasetSpec': {
                        'additionalProperties': false,
                        'description': 'DTO specifying a structured JSONL dataset to evaluate',
                        'properties': {
                            'input_file': {
                                'description': 'Path to a local JSONL-file containing the dataset to evaluate.',
                                'format': 'file-path',
                                'title': 'Input File',
                                'type': 'string'
                            },
                            'record_id_key': {
                                'default': '#IDX#',
                                'description': '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.',
                                'minLength': 1,
                                'title': 'Record Id Key',
                                'type': 'string'
                            },
                            'text_key': {
                                'description': 'The dataset field/ column name that contains the text to evaluate, specified as a non-empty string.',
                                'minLength': 1,
                                'title': 'Text Key',
                                'type': 'string'
                            }
                        },
                        'required': ['input_file', 'text_key'],
                        'title': 'JsonlDatasetSpec',
                        'type': 'object'
                    },
                    'MetricSpec': {
                        'additionalProperties': false,
                        'description': 'DTO specifying a metric and its configuration',
                        'properties': {
                            'name': {
                                'description': 'Non-empty name under which a metric is registered.  !!! note "See also" - [`BaseMetric.name`](https://psandhaas.github.io/evaLLM/reference/evallm/metrics/base/#evallm.metrics.base.BaseMetric.name) - [`@register_metric`](https://psandhaas.github.io/evaLLM/reference/evallm/metrics/registry/#evallm.metrics.registry.register_metric) - [`MetricsCatalog.available`](https://psandhaas.github.io/evaLLM/reference/evallm/metrics/registry/#evallm.metrics.registry.MetricsCatalog.available)',
                                'minLength': 1,
                                'title': 'Name',
                                'type': 'string'
                            },
                            'config': {
                                'anyOf': [
                                    {'additionalProperties': {'$ref': '#/$defs/JsonValue'}, 'type': 'object'},
                                    {'type': 'null'}
                                ],
                                'default': null,
                                'description': 'Runtime configuration for the metric or `null` (default), if not needed.  !!! note "See also" - [`BaseMetric.config_schema`](https://psandhaas.github.io/evaLLM/reference/evallm/metrics/base/#evallm.metrics.base.BaseMetric.config_schema) - [`BaseMetric.from_config`](https://psandhaas.github.io/evaLLM/reference/evallm/metrics/base/#evallm.metrics.base.BaseMetric.from_config)',
                                'title': 'Config'
                            }
                        },
                        'required': ['name'],
                        'title': 'MetricSpec',
                        'type': 'object'}
                    },
         'additionalProperties': false,
         'description': 'DTO specifying which configuration of metrics to evaluate a given dataset on',
         'properties': {
            'dataset': {
                '$ref': '#/$defs/JsonlDatasetSpec',
                'description': 'The dataset to evaluate, specified as a [`JsonlDatasetSpec`](https://psandhaas.github.io/evaLLM/reference/evallm/application/models/#evallm.application.models.JsonlDatasetSpec).'
            },
            'metrics': {
                'description': 'List of one or more metrics to use for evaluation, each specified as a [`MetricSpec`](https://psandhaas.github.io/evaLLM/reference/evallm/application/models/#evallm.application.models.MetricSpec).',
                'items': {'$ref': '#/$defs/MetricSpec'},
                'minItems': 1,
                'title': 'Metrics',
                'type': 'array'
            }
         },
         'required': ['dataset', 'metrics'],
         'title': 'EvaluationRequest',
         'type': 'object'}
        ```
    """

    _resolved_records: AbstractContextManager[Iterator[DatasetRecord]]
    _composite_metric: CompositeMetric

    def __init__(
        self, request: EvaluationRequest, *, catalog: type[Catalog] = MetricsCatalog
    ) -> None:
        self._resolved_records, self._composite_metric = RequestResolver(
            catalog=catalog
        )(request)

    def _evaluate_record(self, record: DatasetRecord) -> dict[str, JsonValue]:
        """Evaluate a single dataset record using the composite metric.

        Args:
            record (DatasetRecord): The dataset record to evaluate.

        Returns:
            dict[str, JsonValue]: A dictionary mapping metric result keys to the values
                that were computed for the record.
        """
        composite_name: str = self._composite_metric.name
        return cast(
            dict[str, JsonValue],
            self._composite_metric.evaluate(record.text)[composite_name],
        )

    def run(self) -> Iterator[EvaluationResult]:
        """Evaluate the records in the request using the specified metrics.

        This method iterates over the resolved dataset records and applies the
        composite metric to each record, yielding an
        [`EvaluationResult`][evallm.application.models.EvaluationResult] for each one.

        Yields:
            EvaluationResult: An evaluation result for each record, containing the
                record's unique identifier and a dictionary mapping metric result keys
                to their computed values.
        """
        with self._resolved_records as records:
            for record in records:
                yield EvaluationResult(
                    record_id=record.record_id,
                    results=self._evaluate_record(record),
                )
