import pytest
from pr.core.context import should_compress_context, compress_context
from pr.config import RECENT_MESSAGES_TO_KEEP


class TestContextManagement:

    def test_should_compress_context_below_threshold(self):
        messages = [{'role': 'user', 'content': 'test'}] * 10
        assert should_compress_context(messages) is False

    def test_should_compress_context_above_threshold(self):
        messages = [{'role': 'user', 'content': 'test'}] * 35
        assert should_compress_context(messages) is True

    def test_compress_context_preserves_system_message(self):
        messages = [
            {'role': 'system', 'content': 'System prompt'},
            {'role': 'user', 'content': 'Hello'},
            {'role': 'assistant', 'content': 'Hi'},
        ] * 40  # Ensure compression
        compressed = compress_context(messages)
        assert compressed[0]['role'] == 'system'
        assert 'System prompt' in compressed[0]['content']

    def test_compress_context_keeps_recent_messages(self):
        messages = [{'role': 'user', 'content': f'msg{i}'} for i in range(40)]
        compressed = compress_context(messages)
        # Should keep recent messages
        recent = compressed[-RECENT_MESSAGES_TO_KEEP:]
        assert len(recent) == RECENT_MESSAGES_TO_KEEP
        # Check that the messages are the most recent ones
        for i, msg in enumerate(recent):
            expected_index = 40 - RECENT_MESSAGES_TO_KEEP + i
            assert msg['content'] == f'msg{expected_index}'