#!/usr/bin/env python3

"""![status](../../../badges/status-application-api.svg) Public API, providing an interface for evaluation"""

from __future__ import annotations

import os
from abc import ABC, abstractmethod
from collections.abc import Iterator, Mapping
from pathlib import Path
from typing import Any, Literal, Self, cast

from pydantic import JsonValue

from evallm.application.errors import DuplicateMetricResultKeyError
from evallm.application.models import (
    BaseDatasetSpec,
    EvaluationRequest,
    EvaluationResult,
    JsonlDatasetSpec,
    MetricSpec,
)
from evallm.application.services import Evaluator
from evallm.application.validation import ensure_unique_metric_result_keys
from evallm.metrics import builtin as builtin_metrics  # noqa: F401
from evallm.metrics.registry import Catalog, MetricsCatalog


class BaseEvaluationBuilder[DatasetSpecT: BaseDatasetSpec](ABC):
    """Abstract mutable builder for creating evaluation requests."""

    _catalog: type[Catalog]
    _metrics: list[MetricSpec]

    def __init__(self, *, catalog: type[Catalog] = MetricsCatalog) -> None:
        if not issubclass(catalog, Catalog):
            raise TypeError(
                f"Provided catalog {catalog.__name__} doesn't implement the Catalog interface"
            )
        self._catalog = catalog
        self._metrics = []

    @property
    def catalog(self) -> type[Catalog]:
        """Immutable catalog used for metric resolution.

        Returns:
            type[Catalog]: Catalog class used for metric resolution in this builder.
        """
        return self._catalog

    @property
    def metrics(self) -> tuple[MetricSpec, ...]:
        """Immutable view of attached metric specs.

        Returns:
            tuple[MetricSpec, ...]: Tuple of currently attached metric specs.
        """
        return tuple(self._metrics)

    def use(self, metric_spec: MetricSpec) -> Self:
        """Attach a single metric spec and return the same builder instance.

        Returns:
            Self: The mutated builder instance.

        Raises:
            TypeError: If `metric_spec` is not a
                [`MetricSpec`][evallm.application.models.MetricSpec] instance.
            DuplicateMetricResultKeyError: If attaching the metric would cause a final
                result-key collision.
        """
        if not isinstance(metric_spec, MetricSpec):
            raise TypeError("metric_spec must be an instance of MetricSpec")

        try:
            ensure_unique_metric_result_keys([*self._metrics, metric_spec])
        except DuplicateMetricResultKeyError as e:
            raise e

        self._metrics.append(metric_spec)
        return self

    @abstractmethod
    def build_dataset(self) -> DatasetSpecT:
        """Build the dataset spec for the concrete builder."""

    def build_request(self) -> EvaluationRequest:
        """Build the canonical evaluation request.

        Returns:
            EvaluationRequest: Canonical request assembled from current builder state.

        Raises:
            ValueError: If no metrics are attached.
        """
        if not self._metrics:
            raise ValueError("At least one metric must be attached before building")
        return EvaluationRequest(
            dataset=self.build_dataset(), metrics=list(self._metrics)
        )

    def build_evaluator(self) -> Evaluator:
        """Build an evaluator from the current builder state.

        Returns:
            Evaluator: Evaluator configured with the built request and catalog.
        """
        return Evaluator(self.build_request(), catalog=self._catalog)

    def run(self) -> Iterator[EvaluationResult]:
        """Lazily execute evaluation and yield individual results.

        Yields:
            EvaluationResult: Per-record evaluation result.
        """
        yield from self.build_evaluator().run()

    @classmethod
    @abstractmethod
    def from_request(
        cls,
        request: EvaluationRequest,
        *,
        catalog: type[Catalog] = MetricsCatalog,
    ) -> Self:
        """Rebuild a concrete builder from a canonical request."""
        raise NotImplementedError(
            "from_request must be implemented by concrete builder classes"
        )  # pragma: no cover


class JsonlEvaluationBuilder(BaseEvaluationBuilder[JsonlDatasetSpec]):
    """Concrete builder for JSONL-based evaluations."""

    _input_file: Path
    _text_key: str
    _record_id_key: str

    def __init__(
        self,
        input_file: str | os.PathLike[str],
        *,
        text_key: str,
        record_id_key: str = "#IDX#",
        catalog: type[Catalog] = MetricsCatalog,
    ) -> None:
        super().__init__(catalog=catalog)
        self._input_file = Path(input_file)
        self._text_key = text_key
        self._record_id_key = record_id_key

    @property
    def input_file(self) -> Path:
        """Input JSONL dataset path.

        Returns:
            Path: The input JSONL dataset as a normalized path.
        """
        return self._input_file

    @property
    def text_key(self) -> str:
        """Key of the field containing text to evaluate in each JSONL record.

        Returns:
            str: The field's name.
        """
        return self._text_key

    @property
    def record_id_key(self) -> str:
        """Key of the field containing record ids.

        Returns:
            str | Literal["#IDX#"]: The field's name or the reserved literal "#IDX#" if
                record ids aren't present and should instead be auto-generated as
                0-based indices.
        """
        return self._record_id_key

    def build_dataset(self) -> JsonlDatasetSpec:
        """Build a JSONL dataset spec from the builder state.

        Returns:
            JsonlDatasetSpec: JSONL dataset spec built from current builder fields.
        """
        return JsonlDatasetSpec(
            input_file=self._input_file,
            text_key=self._text_key,
            record_id_key=self._record_id_key,
        )

    @classmethod
    def from_request(
        cls,
        request: EvaluationRequest,
        *,
        catalog: type[Catalog] = MetricsCatalog,
    ) -> JsonlEvaluationBuilder:
        """Rebuild a JSONL builder from an existing request.

        Returns:
            JsonlEvaluationBuilder: Builder reconstructed from the provided request.

        Raises:
            TypeError: If `request.dataset` is not a
                [`JsonlDatasetSpec`][evallm.application.models.JsonlDatasetSpec].
        """
        if not isinstance(request.dataset, JsonlDatasetSpec):
            raise TypeError(
                "JsonlEvaluationBuilder requires an EvaluationRequest with JsonlDatasetSpec"
            )

        builder = cls(
            request.dataset.input_file,
            text_key=request.dataset.text_key,
            record_id_key=request.dataset.record_id_key,
            catalog=catalog,
        )
        for spec in request.metrics:
            builder.use(spec)
        return builder


def metric(
    name: str,
    *,
    config: Mapping[str, JsonValue] | None = None,
    result_key: str | None = None,
) -> MetricSpec:
    """Create a generic metric spec using a metric's name.

    !!! warning
        This function **only performs syntactic validation** so as to create canonical
        `MetricSpec` objects. It **does not guarantee the existence** of the metric's
        name in the catalog, **nor does it validate the provided config** against the
        metric's [config schema][evallm.metrics.base.BaseMetric.config_schema].

    !!! info "See also"
        - [`MetricSpec.resolve_metric(catalog)`][evallm.application.models.MetricSpec.resolve_metric]
            for runtime resolution and validation of metric specs against a catalog.

    Args:
        name (str): Name of the metric as registered in the catalog.
        config (Mapping[str, JsonValue] | None, optional): Optional config dict for the
            metric. (see also:
            [`BaseMetric.config_schema`][evallm.metrics.base.BaseMetric.config_schema])
        result_key (str | None, optional): Optional explicit
            [result key][evallm.metrics.base.BaseMetric.result_key] for the metric.

    Returns:
        MetricSpec: Canonical metric specification.
    """
    copied_config: dict[str, JsonValue] | None = (
        dict(config) if config is not None else None
    )
    return MetricSpec(name=name, config=copied_config, result_key=result_key)


def count_characters(*, result_key: str | None = None) -> MetricSpec:
    """Create a metric spec for counting the number of characters in the evaluated text.

    Args:
        result_key (str | None, optional): Optional explicit
            [result key][evallm.metrics.base.BaseMetric.result_key] for the metric.

    Returns:
        MetricSpec: Character-count metric specification.
    """
    return metric("CountsMetric", result_key=result_key)


def count_tokens(
    *,
    segmentation: Literal["whitespace", "wordpunct"] = "whitespace",
    result_key: str | None = None,
) -> MetricSpec:
    """Create a metric spec for counting the number of tokens in the evaluated text.

    Args:
        segmentation (Literal["whitespace", "wordpunct"], optional): Tokenization
            method to use. Defaults to "whitespace". (see also:
            [`CountsConfig`][evallm.metrics.builtin.CountsConfig])
        result_key (str | None, optional): Optional explicit
            [result key][evallm.metrics.base.BaseMetric.result_key] for the metric.

    Returns:
        MetricSpec: Token-count metric specification.
    """
    return metric(
        "CountsMetric",
        config={
            "segments": "tokens",
            "segmentation": segmentation,
        },
        result_key=result_key,
    )


def count_encodings(
    model: str,
    *,
    hf_token: str | None = None,
    result_key: str | None = None,
) -> MetricSpec:
    """Create a metric spec for counting the number of model-specific encodings in a text.

    This metric allows for specifying any
    [:hugging: HuggingFace model](https://huggingface.co/models?library=transformers)
    to use for segmenting the text into encodings.

    !!! note "Counting encodings"
        The `'hf-tokenizers'` extra of evaLLM must be installed to use this metric, as
        it relies on the HuggingFace tokenizers library for encoding text.

        Install it with:
        ```bash
        pip install "evallm-qa[hf-tokenizers]"
        ```

    Args:
        model (str): The name of a
            [HuggingFace model](https://huggingface.co/models?library=transformers)
            that provides a pretrained tokenizer, which will be used to segment the
            text into model-specific encodings.

    Keyword Args:
        hf_token (str | None): [User access token for HuggingFace](https://huggingface.co/docs/hub/security-tokens).
            This is only required when specifying a
            [gated model](https://huggingface.co/docs/hub/models-gated) to use for
            encoding. Note that you can also set the `HF_TOKEN` environment variable to
            avoid passing the token directly in the configuration.
        result_key (str | None, optional): Optional explicit
            [result key][evallm.metrics.base.BaseMetric.result_key] for the metric.

    Returns:
        MetricSpec: Encoding-count metric specification.
    """
    config: dict[str, JsonValue] = {
        "segments": "encodings",
        "segmentation": model,
    }
    if hf_token is not None:
        config["hf_token"] = hf_token
    return metric("CountsMetric", config=config, result_key=result_key)


def json_format(
    expected_schema: Mapping[str, Any] | str,
    *,
    check_formats: bool = False,
    result_key: str | None = None,
) -> MetricSpec:
    """Create a metric spec for validating JSON output format.

    Returns:
        MetricSpec: JSON-format validation metric specification.
    """
    schema_value: JsonValue = (
        cast(JsonValue, dict(expected_schema))
        if isinstance(expected_schema, Mapping)
        else cast(JsonValue, expected_schema)
    )
    return metric(
        "JsonFormatMetric",
        config={
            "expected_schema": schema_value,
            "check_formats": check_formats,
        },
        result_key=result_key,
    )


__all__ = [
    "BaseEvaluationBuilder",
    "JsonlEvaluationBuilder",
    "count_characters",
    "count_encodings",
    "count_tokens",
    "json_format",
    "metric",
]
