import unittest
from array import array

from sglang.srt.managers.schedule_batch import Req
from sglang.srt.sampling.sampling_params import SamplingParams


STOP_ID = 1
EOS_ID = 2
ID_TO_TEXT = {
    STOP_ID: "STOP",
    EOS_ID: "",
    **{i: chr(ord("a") + i % 26) for i in range(10, 40)},
}


class _FakeTokenizer:
    eos_token_id = EOS_ID
    additional_stop_token_ids = None

    def decode(self, ids):
        return "".join(ID_TO_TEXT[int(i)] for i in ids)


class _NormalizeTokenizer:
    def encode(self, text, add_special_tokens=False):
        return list(range(len(text)))


def _make_req(output_ids, *, stop=None, eos_token_ids=frozenset()):
    sampling_params = SamplingParams(max_new_tokens=1000, stop=stop)
    sampling_params.normalize(tokenizer=_NormalizeTokenizer())
    req = Req(
        rid="r",
        origin_input_text="",
        origin_input_ids=array("q", [0]),
        sampling_params=sampling_params,
        eos_token_ids=set(eos_token_ids),
        vocab_size=10_000,
    )
    req.tokenizer = _FakeTokenizer()
    req.output_ids = array("q", output_ids)
    return req


class TestStopStrSpeculative(unittest.TestCase):
    def test_stop_str_wins_over_eos_in_same_spec_step(self):
        # A speculative step may accept both the stop string and EOS. The stop
        # string must finish first so finished_len trims at STOP instead of EOS.
        req = _make_req([10, 11, STOP_ID, EOS_ID], stop=["STOP"], eos_token_ids={EOS_ID})

        req.check_finished(new_accepted_len=4)

        self.assertTrue(req.finished())
        self.assertEqual(req.finished_reason.matched, "STOP")
        self.assertEqual(req.finished_len, 3)
        self.assertEqual(list(req.output_ids_through_stop), [10, 11, STOP_ID])


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