from ctransformers import LLM, Config


class MockLLM(LLM):
    def __init__(self):
        self._config = Config()
        self._llm = None

    @property
    def config(self):
        return self._config

    def tokenize(self, prompt, **kwargs):
        self.tokens = prompt.split(" ")
        return range(len(self.tokens))

    def generate(self, tokens, **kwargs):
        return tokens

    def detokenize(self, tokens, decode=True):
        text = " " + self.tokens[tokens[0]]
        if not decode:
            text = text.encode()
        return text


class TestLLM:
    def test_stop(self):
        llm = MockLLM()
        prompt = "foo bar baz lorem ipsum\ndolor"
        expected = [
            ([], " foo bar baz lorem ipsum\ndolor"),
            (["dolor "], " foo bar baz lorem ipsum\ndolor"),
            (["ipsum "], " foo bar baz lorem ipsum\ndolor"),
            (["doloro"], " foo bar baz lorem ipsum\ndolor"),
            (["ipsumo"], " foo bar baz lorem ipsum\ndolor"),
            (["dolor"], " foo bar baz lorem ipsum\n"),
            (["ipsum"], " foo bar baz lorem "),
            (["olor"], " foo bar baz lorem ipsum\nd"),
            (["olo"], " foo bar baz lorem ipsum\nd"),
            (["psum"], " foo bar baz lorem i"),
            (["psu"], " foo bar baz lorem i"),
            (["z lor"], " foo bar ba"),
            (["rem", "or"], " foo bar baz l"),
            (["foo"], " "),
            (["f"], " "),
            ([" "], ""),
            (["\n"], " foo bar baz lorem ipsum"),
            (["m\nd"], " foo bar baz lorem ipsu"),
        ]
        for stop, response in expected:
            assert llm(prompt, stop=stop) == response
            if len(stop) == 1:
                assert llm(prompt, stop=stop[0]) == response