import unittest
from types import SimpleNamespace
import mlx.core as mx
from generate import generate


class Runner(unittest.TestCase):
    def model(self):
        class Tokenizer:
            def encode(self,*args,**kwargs):return SimpleNamespace(ids=[7,8])
            def decode(self,ids,**kwargs):return ' '.join(str(i) for i in ids if i)
        class Model:
            max_tokens=128
            w=SimpleNamespace(config={'eos_token_id':0})
            hash=SimpleNamespace(history=[])
            tokenizer=Tokenizer()
            def step(self,token):
                self.position+=1
                expected=20 if self.position<=2 else 21 if self.position==3 else 0
                return mx.array([10. if i==expected else -10. for i in range(32)]),mx.array([0.])
        return Model()

    def test_greedy_output_counts_and_context_reset(self):
        model=self.model()
        for _ in range(2):
            output,report=generate(model,'Hello',16)
            self.assertEqual(output,'20 21')
            self.assertEqual(report['output_tokens'],3)
            self.assertEqual(report['timed_decode_steps'],2)
            self.assertEqual(report['stop_reason'],'eos')

    def test_output_limit(self):
        output,report=generate(self.model(),'Hello',1)
        self.assertEqual(output,'20');self.assertEqual(report['timed_decode_steps'],0)

    def test_invalid_budget(self):
        with self.assertRaises(ValueError):generate(self.model(),'Hello',0)


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