# Copyright 2026. Released under the MIT license.
"""Character/control-token tokenizer for the 3-digit-basic-calc calculator model.

Digits, operators and symbols are single characters; the scratchpad grammar adds
multi-character control tokens (``<add>``, ``<step>``, ``<qmul>`` ...). A minus
sign is emitted as the unary token ``~`` at the start of an expression, after
``=``, or immediately after a binary operator, so ``12*-34`` and ``5--3`` stay
unambiguous while ``-`` remains the binary subtraction operator.
"""

from __future__ import annotations

import json
import os

from transformers import PreTrainedTokenizer


SPECIAL_TOKENS = ["<pad>", "<bos>", "<eos>", "<unk>"]
DIGITS = list("0123456789")
OPERATORS = ["+", "-", "/", "*"]
SYMBOLS = ["=", ".", "~"]
CONTROL_TOKENS = [
    "<add>", "<sub>", "<mul>", "<div>", "<pos>", "<neg>",
    "<state>", "<step>", "<ans>", "<nan>", "<col>", "<qmul>", "<rem>",
]
TOKENS = SPECIAL_TOKENS + DIGITS + OPERATORS + SYMBOLS + CONTROL_TOKENS
_UNARY_PREDECESSORS = {"=", "+", "-", "/", "*"}


class ThreeDigitBasicCalcTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "vocab.json"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(self, vocab_file=None, **kwargs):
        if vocab_file and os.path.isfile(vocab_file):
            with open(vocab_file, encoding="utf-8") as handle:
                self._vocab = json.load(handle)
        else:
            self._vocab = {tok: i for i, tok in enumerate(TOKENS)}
        self._ids_to_tokens = {i: t for t, i in self._vocab.items()}
        self._ordered_controls = sorted(CONTROL_TOKENS, key=len, reverse=True)
        kwargs.setdefault("pad_token", "<pad>")
        kwargs.setdefault("bos_token", "<bos>")
        kwargs.setdefault("eos_token", "<eos>")
        kwargs.setdefault("unk_token", "<unk>")
        # Generation takes the logits from the final sequence position.  Left
        # padding therefore keeps the final position on a real prompt token for
        # every member of a mixed-length batch.
        kwargs.setdefault("padding_side", "left")
        super().__init__(**kwargs)

    @property
    def vocab_size(self):
        return len(self._vocab)

    def get_vocab(self):
        return dict(self._vocab)

    def _tokenize(self, text):
        tokens = []
        index = 0
        while index < len(text):
            control = next(
                (t for t in self._ordered_controls if text.startswith(t, index)),
                None,
            )
            if control is not None:
                tokens.append(control)
                index += len(control)
                continue
            ch = text[index]
            if ch == "-" and (
                    index == 0 or text[index - 1] in _UNARY_PREDECESSORS):
                tokens.append("~")
            else:
                tokens.append(ch)
            index += 1
        return tokens

    def _convert_token_to_id(self, token):
        return self._vocab.get(token, self._vocab["<unk>"])

    def _convert_id_to_token(self, index):
        return self._ids_to_tokens.get(index, "<unk>")

    def convert_tokens_to_string(self, tokens):
        out = []
        for tok in tokens:
            if tok in SPECIAL_TOKENS:
                continue
            out.append("-" if tok == "~" else tok)
        return "".join(out)

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        bos = [self._vocab["<bos>"]]
        if token_ids_1 is None:
            return bos + token_ids_0
        return bos + token_ids_0 + token_ids_1

    def save_vocabulary(self, save_directory, filename_prefix=None):
        path = os.path.join(
            save_directory,
            (filename_prefix + "-" if filename_prefix else "") + "vocab.json",
        )
        with open(path, "w", encoding="utf-8") as handle:
            json.dump(self._vocab, handle, ensure_ascii=False, indent=2)
        return (path,)