#!/usr/bin/env python3

"""`CountsMetric` — count the number of segments in an input text

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

from __future__ import annotations

__all__ = []

import os
import re
from collections.abc import Callable, Sequence
from functools import cached_property
from typing import Any, Literal, Self

import orjson
from pydantic import (
    BaseModel,
    Field,
    ModelWrapValidatorHandler,
    computed_field,
    model_validator,
)

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

# HuggingFace only available if extras are installed
try:
    from huggingface_hub import auth_check, file_exists
    from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
    from tokenizers import Tokenizer

    HF_AVAILABLE = True
except ImportError:
    HF_AVAILABLE = False


class CountsConfig(BaseModel):
    """The configuration schema for the [`CountsMetric`][evallm.metrics.builtin.counts.CountsMetric].

    Keyword Args:
        segments (Literal["characters", "tokens", "encodings"] | None): (default
            = `"characters"`) Which kind of segments to count. If `"characters"`, the
            text is stripped of leading & trailing whitespace and each remaining
            character is considered a segment. See the `segmentation` parameter for how
            the text is split into segments when `segments` is set to `"tokens"` or
            `"encodings"`.
        segmentation (Literal["whitespace", "wordpunct"] | str | None): How to split
            the input text into segments.

            If `segments = "characters"`, this parameter is ignored.

            If `segments = "tokens"`, this *can* be

            - "whitespace" to split on any whitespace (default), or
            - "wordpunct" to tokenize the text into words and punctuation using
                `"\\w+|[^\\w\\s]+"` as a regular expression.

            If `segments = "encodings"`, this *must* be 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.
        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.


    ??? example annotate "Example configuration"
        ```python
        characters: CountsMetric = CountsMetric()  # (1)!

        tokens_ws: CountsMetric = CountsMetric.from_config(
            {
                "segments": "tokens"  # (2)!
            }
        )

        tokens_wordpunct: CountsMetric = CountsMetric.from_config(
            **{
                "segments": "tokens",
                "segmentation": "wordpunct",  # (3)!
            }
        )

        encs_bart: CountsMetric = CountsMetric.from_config(
            {
                "segments": "encodings",  # (4)!
                "segmentation": "facebook/bart-base",
            }
        )


        encs_bert_amharic: CountsMetric = CountsMetric.from_config(
            **{
                "segments": "encodings",
                "segmentation": "rasyosef/bert-amharic-tokenizer",  # (5)!
            }
        )

        encs_llama: CountsMetric = CountsMetric.from_config(
            **{
                "segments": "encodings",
                "segmentation": "meta-llama/Llama-3.3-70B-Instruct",
                "hf_token": "hf_...",  # (6)!
            }
        )
        ```

    1.  By default, the `CountsMetric` simply counts the number of characters.
    2.  To count tokens instead of characters, set `segments="tokens"`. This will split
        the text into tokens separated by one or more whitespace characters by default.
    3.  Use a slightly more sophisticated tokenization scheme.
    4.  To count model-specific encodings, set `segments="encodings"` and specify the
        model to use for encoding in the `segmentation` field.
    5.  Maybe you're working with non-English texts? No problem, just specify a model
        that provides a tokenizer pretrained on the relevant language.
    6.  If you want to use a gated model for encoding, make sure to provide a valid
        HuggingFace user access token with the necessary permissions to access the
        model. You can also set the `HF_TOKEN` environment variable instead if you
        prefer not to pass the token in directly.
    """

    segments: Literal["characters", "tokens", "encodings"] = Field(
        default="characters",
        description=str(
            "Which kind of segments to count. Options are: "
            + "`'characters'` (default), `'tokens'`, or `'encodings'`."
        ),
    )
    segmentation: str | None = Field(
        default=None,
        description=str(
            "How to split the input text into segments. If `segments = 'characters'`, "
            + "this parameter is ignored. If `segments = 'tokens'`, this can be "
            + "`'whitespace'` to split on any whitespace (default), or `'wordpunct'` "
            + "to tokenize the text into words and punctuation using `\\w+|[^\\w\\s]+` "
            + "as a regular expression. If `segments = 'encodings'`, this must be the "
            + "name of a HuggingFace model that provides a pretrained tokenizer, which "
            + "will be used to segment the text into model-specific encodings."
        ),
    )
    hf_token: str | None = Field(
        default=None,
        description=str(
            "User access token for HuggingFace. This is only required when specifying a "
            + "gated model to use for encoding. Ignored if the `HF_TOKEN` "
            + "environment variable is set."
        ),
        exclude=True,
    )

    @computed_field  # type: ignore[prop-decorator]
    @cached_property
    def tokenizer(self) -> Callable[[str], Sequence[str | int]]:
        """The function to use for segmenting the input text, based on the
        configuration.

        Returns:
            Callable[[str], Sequence[str | int]]: A function that takes a string as
                input and returns a sequence of segments, where each segment is either
                a string (for tokenization) or an integer (for encoding).
        """
        if self.segments == "characters":
            return lambda text: list(text.strip())
        elif self.segments == "tokens":
            if self.segmentation == "wordpunct":
                return lambda text: re.findall(r"\w+|[^\w\s]+", text)
            else:  # default to whitespace tokenization
                return lambda text: text.split()
        else:
            tokenizer: Tokenizer = Tokenizer.from_pretrained(self.segmentation)
            return lambda text: tokenizer.encode(text).ids

    @model_validator(mode="wrap")
    @classmethod
    def _validate_config(
        cls, config: str | dict[str, Any], handler: ModelWrapValidatorHandler[Self]
    ) -> Self:
        """Validate the configuration for the `CountsMetric`.

        This method performs the following validations:

        - If `segments` is set to `"tokens"`, it checks that `segmentation` is either
          `"whitespace"` or `"wordpunct"`.
        - If `segments` is set to `"encodings"`, it checks that `segmentation` is a
            valid HuggingFace model name and - if the model is gated - that a valid
            HuggingFace user access token is available as an environment variable or in
            the configuration.
        - If an accessible model is specified, it checks that the model provides a
            pretrained tokenizer that can be loaded successfully.

        Args:
            config (str | dict[str, Any]): The configuration to validate, either as a
                JSON string or a dictionary.
            handler (ModelWrapValidatorHandler[Self]): The Pydantic model wrap validator
                handler, used to construct the model instance after validation.

        Returns:
            CountsConfig: The validated configuration as a `CountsConfig` instance.

        Raises:
            ValueError: If the validation fails for any of the reasons described above
                or because a string payload can't be parsed into a dictionary.
        """
        if isinstance(config, str):
            try:
                loaded = orjson.loads(config)
            except orjson.JSONDecodeError as e:
                raise ValueError(
                    "Invalid JSON-string. Failed to parse configuration."
                ) from e
            if not isinstance(loaded, dict):  # pragma: no cover
                raise ValueError("Configuration JSON must decode to an object.")
            config = loaded

        if (segments := config.get("segments")) == "encodings":
            if (segmentation := config.get("segmentation")) is None:
                raise ValueError(
                    "When segments='encodings', 'segmentation' must be specified as "
                    + "the name of a HuggingFace model."
                )
            if not HF_AVAILABLE:
                raise ValueError(
                    "HuggingFace models require optional dependencies. Please install "
                    + "evaLLM with the 'hf-tokenizers' extra to use this feature."
                )

            # env variable takes precedence over config value for token
            token: str | Literal[False] = os.getenv(
                "HF_TOKEN", config.get("hf_token", False)
            )
            try:
                auth_check(segmentation, token=token)
            except GatedRepoError as e:
                raise ValueError(
                    "You do not have permission to access this gated repository."
                ) from e
            except RepositoryNotFoundError as e:
                raise ValueError(
                    "The repository was not found or you do not have access."
                ) from e

            # check if the model provides a pretrained tokenizer
            if not file_exists(segmentation, "tokenizer.json", token=token):
                raise ValueError(
                    f"The specified model '{segmentation}' does not provide a "
                    + "pretrained tokenizer that can be loaded with the HuggingFace "
                    + "Tokenizers library."
                )

        elif segments == "tokens":
            if (segmentation := config.get("segmentation", "whitespace")) not in (
                "whitespace",
                "wordpunct",
            ):
                raise ValueError(
                    f"Unknown segmentation method '{segmentation}'. Specify either "
                    + "'whitespace' (default) or 'wordpunct'."
                )
        elif segments is not None and segments != "characters":
            raise ValueError(
                f"Unknown segment type '{segments}'. Specify either 'characters' "
                + "(default), 'tokens', or 'encodings'."
            )
        else:
            segments = "characters"  # default

        return handler(config)


@register_metric
@config_schema(schema=CountsConfig)
class CountsMetric(BaseMetric[int]):
    """Metric for counting segments in an input text.

    !!! info "Evaluation results"
        Given an input string, this metric produces a dictionary mapping the metric's
        name to a single count. The count represents the number of segments that are
        contained in the input text, where the type of segments counted (characters,
        tokens, or encodings) can be configured through the `segments` and
        `segmentation` runtime parameters. If no configuration is specified, this
        metric will simply count the number of characters in the input text (after
        stripping leading and trailing whitespace) by default.

        Note that specifiying a model to use for encoding input texts requires
        installing evaLLM with the `hf-tokenizers` extra:
        ```bash
        pip install evallm-qa[hf-tokenizers]
        ```

    Keyword Args:
        segments (Literal["characters", "tokens", "encodings"] | None): (default
            = `"characters"`) Which kind of segments to count. If `"characters"`, the
            text is stripped of leading & trailing whitespace and each remaining
            character is considered a segment. See the `segmentation` parameter for how
            the text is split into segments when `segments` is set to `"tokens"` or
            `"encodings"`.
        segmentation (Literal["whitespace", "wordpunct"] | str | None): How to split
            the input text into segments.

            If `segments = "characters"`, this parameter is ignored.

            If `segments = "tokens"`, this *can* be

            - "whitespace" to split on any whitespace (default), or
            - "wordpunct" to tokenize the text into words and punctuation using
                `"\\w+|[^\\w\\s]+"` as a regular expression.

            If `segments = "encodings"`, this *must* be 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.
        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.
    """

    def evaluate(self, text: str) -> dict[str, int]:
        """Evaluate how many segments are contained in the input text.

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

        Returns:
            dict[str, int]: A dictionary mapping this metric instance's
                [`result_key`][evallm.metrics.base.BaseMetric.result_key] to the count
                of segments in the input text. The type of segments counted
                (characters, tokens, or encodings) is determined by the `segments`
                configuration parameter, and how the text is split into segments is
                determined by the `segmentation` configuration parameter when
                applicable.
        """
        if text is None or text.strip() == "":
            res: int = 0
        else:
            res = len(self.config.get("tokenizer")(text))  # type: ignore[misc]

        return self.result(res)
