"""Tests for prompt_builder.tokens"""

from prompt_builder import TokenCounter, get_counter


class TestTokenCounter:
    def test_count_empty(self):
        tc = TokenCounter()
        assert tc.count("") == 0

    def test_count_simple(self):
        tc = TokenCounter()
        tokens = tc.count("Hello, world!")
        assert tokens > 0
        assert tokens < 10  # sanity check

    def test_count_longer_text(self):
        tc = TokenCounter()
        short = tc.count("Hi")
        long = tc.count("This is a much longer sentence with many more words in it.")
        assert long > short

    def test_count_messages(self):
        tc = TokenCounter()
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hello"},
        ]
        total = tc.count_messages(messages)
        # Should include content tokens + overhead
        assert total > tc.count("You are helpful.") + tc.count("Hello")

    def test_count_messages_empty(self):
        tc = TokenCounter()
        total = tc.count_messages([])
        assert total == 2  # just the reply priming

    def test_fallback_encoding(self):
        """Unknown model should fallback to cl100k_base."""
        tc = TokenCounter(model="some-unknown-model-xyz")
        tokens = tc.count("Hello")
        assert tokens > 0

    def test_get_counter_reuse(self):
        c1 = get_counter("gpt-4")
        c2 = get_counter("gpt-4")
        assert c1 is c2

    def test_get_counter_different_model(self):
        c1 = get_counter("gpt-4")
        c2 = get_counter("gpt-3.5-turbo")
        # After getting c2, the module default is now gpt-3.5-turbo
        assert c2.model == "gpt-3.5-turbo"
