"""Module for tokenization using regex patterns.

Functions returning regex patterns are generally advised to return uncompiled patterns to allow combination, but
compiled patterns are also allowed, which may be necessary when passing flags to the regex compiler.
"""

from typing import Union

import regex as re

__all__ = ["Tokenizer", "whitespace_pattern", "strip_punctuation_pattern", "strip_punctuation_keep_symbols_pattern"]


def whitespace_pattern() -> re.Pattern:
    """Return a regex pattern that matches sequences of non-whitespace characters.

    Returns:
        re.Pattern: The compiled regex pattern.
    """
    return re.compile(r"\S+")


def strip_punctuation_pattern(split_on_escaped: str | None = None, split_on_pattern: str | None = None) -> re.Pattern:
    """Return a regex pattern that matches tokens without internal whitespace or punctuation per specified characters.

    Args:
        split_on_escaped (str): A string of characters to split on in addition to whitespace. Will be escaped.
        split_on_pattern (str): A regex pattern to split on in addition to whitespace. Will be used as-is.

    Returns:
        re.Pattern: The compiled regex pattern.
    """
    letters_digits = r"\p{L}\p{N}"
    symbols_punctuation_marks = r"\p{S}\p{P}\p{M}"
    if split_on_escaped is None and split_on_pattern is None:
        return re.compile(rf"[{letters_digits}]+([[{symbols_punctuation_marks}]]+[{letters_digits}]+)*", re.V1)
    escaped_split_on = (re.escape(split_on_escaped) if split_on_escaped is not None else "") + (split_on_pattern or "")
    return re.compile(
        rf"[{letters_digits}]+([[{symbols_punctuation_marks}]--[{escaped_split_on}]]+[{letters_digits}]+)*", re.V1
    )


def strip_punctuation_keep_symbols_pattern(
    split_on_escaped: str | None = None,
) -> re.Pattern:
    """
    Return a regex pattern that matches tokens without internal whitespace or punctuation per specified characters, but
    keeps currency symbols, math symbols, and percent signs as separate tokens.

    Args:
        split_on_escaped (str): A string of characters to split on in addition to whitespace. Will be escaped.

    Returns:
        re.Pattern: The compiled regex pattern.
    """
    math_currency = r"\p{Sm}\p{Sc}%"
    strip_pattern = strip_punctuation_pattern(
        split_on_escaped=split_on_escaped,
        split_on_pattern=math_currency,
    )
    return re.compile(rf"([{math_currency}])|({strip_pattern.pattern})", re.V1)


class Tokenizer(object):
    def __init__(
        self,
        pattern: Union[str, re.Pattern],
        name: str = "__no_name__",
    ):
        """Initialize a tokenizer with a regex pattern.

        Args:
            name (str): The config name of the tokenizer.
            pattern (Union[str, re.Pattern]): The regex pattern to use for tokenization.
        """
        self._pattern = re.compile(pattern) if isinstance(pattern, str) else pattern
        self._name = name

    @property
    def pattern(self) -> re.Pattern:
        """Get the regex pattern used by the tokenizer."""
        return self._pattern

    def __call__(self, text: str) -> list[re.Match]:
        """Tokenize the input text using the regex pattern.

        Args:
            text: The input text to tokenize.

        Returns:
            list[re.Match]: A list of regex match objects.
        """
        return list(self._pattern.finditer(text))

    def __repr__(self):
        """Return a string representation of the Tokenizer object."""
        if self._name is None:
            return f"Tokenizer({self._pattern.pattern})"
        return f"Tokenizer({self._name}: {self._pattern.pattern})"
