import torch

from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
from sglang.srt.speculative.spec_utils import traverse_tree


class _TypeCheckingGrammar(BaseGrammarObject):
    def __init__(self):
        super().__init__()
        self.accepted = []
        self.rollback_calls = []

    def accept_token(self, token: int) -> None:
        if not isinstance(token, int):
            raise TypeError(f"expected Python int token, got {type(token)!r}")
        self.accepted.append(token)

    def rollback(self, k: int):
        self.rollback_calls.append(k)

    def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
        vocab_mask[idx].fill_(-1)

    def allocate_vocab_mask(self, vocab_size: int, batch_size: int, device):
        return torch.zeros((batch_size, (vocab_size + 31) // 32), dtype=torch.int32)


def test_traverse_tree_passes_python_int_tokens_to_grammar():
    grammar = _TypeCheckingGrammar()
    retrieve_next_token = torch.tensor([1, -1], dtype=torch.int32)
    retrieve_next_sibling = torch.tensor([-1, -1], dtype=torch.int32)
    draft_tokens = torch.tensor([0, 5], dtype=torch.int64)
    allocate_token_bitmask = torch.zeros((2, 2), dtype=torch.int32)

    traverse_tree(
        retrieve_next_token,
        retrieve_next_sibling,
        draft_tokens,
        grammar,
        allocate_token_bitmask,
        vocab_size=64,
    )

    assert grammar.accepted == [5]
    assert grammar.rollback_calls == [1]