"""TinyBrainBot units-first number tokenizer (byte-level BPE + length-prefixed, reversed digit runs).

Every run of 1-8 digits is written as a length token and then its digits units first:
    "62"   -> "<len2>26"        "1234" -> "<len4>4321"
so a model emits the units digit first and carries flow forward in time. Digit runs touching ":" (clock
times like 3:15) and runs longer than 8 digits are left unchanged. The text is rewritten on encode and
restored on decode, so callers see ordinary numbers. The rewrite is idempotent (an already-rewritten run
is not rewritten again), so it is safe on every encode path.

The idea (units-first digit tokens; a length prefix) follows Atom2.7m and Pascaline-3M. This is our own
implementation.
"""
import re

from transformers import PreTrainedTokenizerFast

MAX_LEN = 8
# a digit run: not part of "<lenN>", not right after a "<lenN>" tag, not touching another digit or ":"
_RUN = re.compile(r"(?<!<len)(?<!<len\d>)(?<![\d:])(\d{1,%d})(?![\d:])" % MAX_LEN)
_TAG = re.compile(r"<len([1-%d])>(\d+)" % MAX_LEN)


def encode_numbers(text):
    if not isinstance(text, str):
        return text
    return _RUN.sub(lambda m: "<len%d>%s" % (len(m.group(1)), m.group(1)[::-1]), text)


def decode_numbers(text):
    def back(m):
        n = int(m.group(1)); d = m.group(2)
        return d[:n][::-1] + d[n:]
    return _TAG.sub(back, text)


def _map(x):
    if isinstance(x, str):
        return encode_numbers(x)
    if isinstance(x, tuple):
        return tuple(_map(v) for v in x)
    if isinstance(x, list):
        return [_map(v) for v in x]
    return x


class TBBUnitsFirstTokenizerFast(PreTrainedTokenizerFast):
    def _batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
        if not kwargs.get("is_split_into_words", False):
            batch_text_or_text_pairs = _map(batch_text_or_text_pairs)
        return super()._batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)

    def _encode_plus(self, text, text_pair=None, *args, **kwargs):
        if not kwargs.get("is_split_into_words", False):
            text, text_pair = _map(text), _map(text_pair)
        return super()._encode_plus(text, text_pair, *args, **kwargs)

    def tokenize(self, text, *args, **kwargs):
        return super().tokenize(encode_numbers(text), *args, **kwargs)

    def _decode(self, token_ids, *args, **kwargs):
        return decode_numbers(super()._decode(token_ids, *args, **kwargs))