import unittest

from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject


class _InnerGrammar:
    def __init__(self):
        self.accepted_tokens = []

    def accept_token(self, token: int):
        self.accepted_tokens.append(token)

    def is_terminated(self):
        return False

    def rollback(self, k: int):
        self.accepted_tokens = self.accepted_tokens[:-k]

    def allocate_vocab_mask(self, vocab_size: int, batch_size: int, device):
        return None

    def fill_vocab_mask(self, vocab_mask, idx: int):
        return None

    def move_vocab_mask(self, vocab_mask, device):
        return vocab_mask

    @property
    def apply_vocab_mask(self):
        return None

    def copy(self):
        ret = _InnerGrammar()
        ret.accepted_tokens = list(self.accepted_tokens)
        return ret

    @property
    def finished(self):
        return False


class TestReasonerGrammarCurrentToken(unittest.TestCase):
    def test_wrapper_tracks_current_token_for_disagg_prebuilt_dedup(self):
        inner = _InnerGrammar()
        grammar = ReasonerGrammarObject(inner, think_end_id=7)
        grammar.maybe_init_reasoning(True)

        grammar.accept_token(10)  # still in thinking; inner grammar is untouched
        self.assertEqual(grammar.current_token, 10)
        self.assertEqual(inner.accepted_tokens, [])

        grammar.accept_token(7)  # exits thinking
        self.assertEqual(grammar.current_token, 7)
        self.assertEqual(inner.accepted_tokens, [])

        grammar.accept_token(58)  # generation token accepted by inner grammar
        self.assertEqual(grammar.current_token, 58)
        self.assertEqual(inner.accepted_tokens, [58])

        # Mirrors disaggregation/decode_schedule_batch_mixin.py:process_prebuilt.
        # Once current_token is tracked on the wrapper, a re-prebuilt request
        # skips re-accepting the already accepted token.
        if grammar.current_token is None:
            grammar.accept_token(58)
        self.assertEqual(inner.accepted_tokens, [58])


if __name__ == "__main__":
    unittest.main()