#!/usr/bin/env python3

"""![status](../../../badges/status-application-adapters.svg) Internal adapters for application-layer DTOs"""

from __future__ import annotations

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

from pydantic import JsonValue, ValidationError

from evallm.application.errors import MissingRecordIdFieldError, MissingTextFieldError
from evallm.application.models import DatasetRecord, JsonlDatasetSpec
from evallm.readers import BaseReader


@contextmanager
def resolved_records(
    reader: BaseReader[dict[str, Any]],
    *,
    spec: JsonlDatasetSpec,
) -> Iterator[Iterator[DatasetRecord]]:
    """Context manager that resolves an evaluation dataset specification to an iterator
    over `DatasetRecord`s.

    Args:
        reader (BaseReader[dict[str, Any]]): Reader yielding raw dataset records as
            dictionaries.
        spec (JsonlDatasetSpec): Dataset specification containing
            [`record_id_key`][evallm.application.models.JsonlDatasetSpec.record_id_key]
            and [`text_key`][evallm.application.models.JsonlDatasetSpec.text_key].

    Yields:
        Iterator[DatasetRecord]: An iterator over lazily resolved
            [`DatasetRecord`][evallm.application.models.DatasetRecord]s.

    Raises:
        MissingRecordIdFieldError: Raised during iteration if a record is missing
            the specified `record_id_key` field.
        MissingTextFieldError: Raised during iteration if a record is missing
            the configured text field.
        ValueError: Raised during iteration if a raw record cannot be validated
            into `DatasetRecord`.
    """  # noqa: DOC502
    with reader as records:

        def _iter() -> Iterator[DatasetRecord]:
            excluded_keys: set[str]
            for index, record in enumerate(records):
                # assign 0-based indices as IDs by default
                if spec.record_id_key == "#IDX#":
                    record_id = str(index)
                    excluded_keys = {spec.text_key}
                else:  # otherwise, use the specified key to extract record IDs
                    if spec.record_id_key not in record:
                        raise MissingRecordIdFieldError(
                            index=index,
                            record_id_key=spec.record_id_key,
                            record=dict(record),
                        )
                    record_id = str(record[spec.record_id_key])
                    excluded_keys = {spec.record_id_key, spec.text_key}

                if spec.text_key not in record:
                    raise MissingTextFieldError(
                        index=index,
                        text_key=spec.text_key,
                        record=dict(record),
                    )

                text = str(record[spec.text_key])
                metadata: dict[str, JsonValue] | None = cast(
                    dict[str, JsonValue] | None,
                    {k: v for k, v in record.items() if k not in excluded_keys} or None,
                )

                try:
                    yield DatasetRecord(
                        record_id=record_id,
                        text=text,
                        metadata=metadata,
                    )
                except ValidationError as exc:
                    raise ValueError(
                        f"Invalid dataset record at index {index}: {exc}"
                    ) from exc

        # yield the generator so that its lifetime is managed by the reader
        yield _iter()
