headroom/tests/test_transforms/test_progressive_summarizer.py

"""Comprehensive tests for progressive summarization.

These tests verify that ProgressiveSummarizer works correctly with:
- Anchored summaries that track message positions
- Callback pattern for summarization (no internal LLM calls)
- CCR integration for retrieval
- Extractive fallback summarization

CRITICAL: NO MOCKS for core logic. All tests use real implementations.
"""

from __future__ import annotations

from typing import Any

import pytest

from headroom.tokenizer import Tokenizer
from headroom.tokenizers import EstimatingTokenCounter
from headroom.transforms.progressive_summarizer import (
    AnchoredSummary,
    ProgressiveSummarizer,
    SummarizationResult,
    extractive_summarizer,
)

# =============================================================================
# Test Fixtures
# =============================================================================


@pytest.fixture
def tokenizer() -> Tokenizer:
    """Create a tokenizer for testing."""
    return Tokenizer(EstimatingTokenCounter())


@pytest.fixture
def simple_conversation() -> list[dict[str, Any]]:
    """Simple conversation without tool calls."""
    return [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello, how are you?"},
        {"role": "assistant", "content": "I'm doing well, thank you for asking!"},
        {"role": "user", "content": "Can you help me with Python?"},
        {"role": "assistant", "content": "Of course! What would you like to know?"},
        {"role": "user", "content": "How do I read a file?"},
        {
            "role": "assistant",
            "content": "You can use open() to read files. Here's an example: with open('file.txt', 'r') as f: content = f.read()",
        },
    ]


@pytest.fixture
def conversation_with_tools() -> list[dict[str, Any]]:
    """Conversation with tool calls and responses."""
    return [
        {"role": "system", "content": "You are a helpful assistant with tools."},
        {"role": "user", "content": "Search for information about Python."},
        {
            "role": "assistant",
            "content": "I'll search for that.",
            "tool_calls": [
                {
                    "id": "call_1",
                    "type": "function",
                    "function": {"name": "search", "arguments": "{}"},
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_1",
            "content": '{"results": [{"title": "Python Guide", "url": "example.com"}, {"title": "Python Tutorial", "url": "tutorial.com"}]}',
        },
        {"role": "assistant", "content": "Here's what I found about Python programming."},
        {"role": "user", "content": "Thanks! Can you search for more?"},
        {
            "role": "assistant",
            "content": "Sure, searching again for more results.",
            "tool_calls": [
                {
                    "id": "call_2",
                    "type": "function",
                    "function": {"name": "search", "arguments": "{}"},
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_2",
            "content": '{"results": [{"title": "Advanced Python", "status": "found"}, {"error": "Some results failed to load"}]}',
        },
        {"role": "assistant", "content": "Here are more results for you."},
    ]


@pytest.fixture
def long_conversation() -> list[dict[str, Any]]:
    """Long conversation for testing summarization scenarios."""
    messages = [{"role": "system", "content": "You are a helpful assistant."}]

    # Add many turns
    for i in range(20):
        messages.append(
            {"role": "user", "content": f"This is question number {i}. What about topic {i}?"}
        )
        messages.append(
            {
                "role": "assistant",
                "content": f"Here's my detailed response about topic {i}. " * 10
                + f"In summary, topic {i} is interesting.",
            }
        )

    return messages


# =============================================================================
# AnchoredSummary Tests
# =============================================================================


class TestAnchoredSummary:
    """Tests for AnchoredSummary dataclass."""

    def test_compression_ratio_calculation(self) -> None:
        """Test compression ratio is calculated correctly."""
        summary = AnchoredSummary(
            summary_text="Summary",
            start_index=0,
            end_index=5,
            original_message_count=6,
            original_tokens=1000,
            summary_tokens=100,
        )
        assert summary.compression_ratio == 0.1  # 100/1000

    def test_compression_ratio_with_zero_original(self) -> None:
        """Test compression ratio handles zero original tokens."""
        summary = AnchoredSummary(
            summary_text="Summary",
            start_index=0,
            end_index=0,
            original_message_count=1,
            original_tokens=0,
            summary_tokens=10,
        )
        assert summary.compression_ratio == 1.0  # fallback

    def test_tokens_saved(self) -> None:
        """Test tokens_saved calculation."""
        summary = AnchoredSummary(
            summary_text="Summary",
            start_index=0,
            end_index=5,
            original_message_count=6,
            original_tokens=1000,
            summary_tokens=100,
        )
        assert summary.tokens_saved == 900

    def test_tokens_saved_no_negative(self) -> None:
        """Test tokens_saved doesn't go negative."""
        summary = AnchoredSummary(
            summary_text="Long summary that is bigger than original",
            start_index=0,
            end_index=0,
            original_message_count=1,
            original_tokens=10,
            summary_tokens=50,
        )
        assert summary.tokens_saved == 0  # max(0, ...)

    def test_optional_fields(self) -> None:
        """Test optional fields have defaults."""
        summary = AnchoredSummary(
            summary_text="Summary",
            start_index=0,
            end_index=5,
            original_message_count=6,
            original_tokens=1000,
            summary_tokens=100,
        )
        assert summary.cache_hash is None
        assert summary.tool_names == []
        assert summary.created_at > 0


# =============================================================================
# Extractive Summarizer Tests
# =============================================================================


class TestExtractiveSummarizer:
    """Tests for the default extractive summarizer."""

    def test_empty_messages(self) -> None:
        """Test handling of empty message list."""
        result = extractive_summarizer([])
        assert result == "[No messages to summarize]"

    def test_simple_conversation(self, simple_conversation: list[dict[str, Any]]) -> None:
        """Test summarization of simple conversation."""
        # Skip system message, use rest
        result = extractive_summarizer(simple_conversation[1:])
        assert "[Summary of 6 messages]" in result
        assert "user messages" in result
        assert "assistant" in result.lower()

    def test_tool_messages_detection(self, conversation_with_tools: list[dict[str, Any]]) -> None:
        """Test that tool messages are detected and counted."""
        result = extractive_summarizer(conversation_with_tools)
        assert "tool outputs" in result.lower()

    def test_error_detection_in_tools(self) -> None:
        """Test that errors in tool responses are detected."""
        messages = [
            {
                "role": "tool",
                "tool_call_id": "call_1",
                "content": "Error: Connection failed",
            },
        ]
        result = extractive_summarizer(messages)
        assert "with errors" in result

    def test_successful_tools(self) -> None:
        """Test that successful tool responses are marked correctly."""
        messages = [
            {
                "role": "tool",
                "tool_call_id": "call_1",
                "content": '{"status": "success", "data": [1, 2, 3]}',
            },
        ]
        result = extractive_summarizer(messages)
        assert "successful" in result

    def test_long_assistant_content_truncated(self) -> None:
        """Test that long assistant content is truncated."""
        messages = [
            {"role": "assistant", "content": "X" * 200},
        ]
        result = extractive_summarizer(messages)
        assert "..." in result  # Truncation indicator

    def test_context_ignored(self) -> None:
        """Test that context parameter exists but doesn't change output format."""
        messages = [{"role": "user", "content": "Hello"}]
        result1 = extractive_summarizer(messages, context="")
        result2 = extractive_summarizer(messages, context="Some context here")
        # Both should work (context is unused in extractive mode)
        assert "[Summary of 1 messages]" in result1
        assert "[Summary of 1 messages]" in result2


# =============================================================================
# ProgressiveSummarizer Core Tests
# =============================================================================


class TestProgressiveSummarizerInit:
    """Tests for ProgressiveSummarizer initialization."""

    def test_default_init(self) -> None:
        """Test default initialization."""
        summarizer = ProgressiveSummarizer()
        assert summarizer.max_summary_tokens == 500
        assert summarizer.min_messages_to_summarize == 3
        assert summarizer.store_for_retrieval is True
        # Default summarizer is extractive_summarizer
        assert summarizer.summarize_fn is not None

    def test_custom_summarize_fn(self) -> None:
        """Test custom summarization function."""

        def custom_fn(messages: list[dict], context: str = "") -> str:
            return f"Custom: {len(messages)} messages"

        summarizer = ProgressiveSummarizer(summarize_fn=custom_fn)
        result = summarizer.summarize_fn([{"role": "user", "content": "test"}])
        assert "Custom: 1" in result

    def test_custom_config(self) -> None:
        """Test custom configuration."""
        summarizer = ProgressiveSummarizer(
            max_summary_tokens=1000,
            min_messages_to_summarize=5,
            store_for_retrieval=False,
        )
        assert summarizer.max_summary_tokens == 1000
        assert summarizer.min_messages_to_summarize == 5
        assert summarizer.store_for_retrieval is False


# =============================================================================
# Find Candidates Tests
# =============================================================================


class TestFindSummarizationCandidates:
    """Tests for finding candidate message groups."""

    def test_no_protected_all_candidates(self) -> None:
        """All messages are candidates when none protected."""
        summarizer = ProgressiveSummarizer(min_messages_to_summarize=3)
        messages = [
            {"role": "user", "content": "1"},
            {"role": "assistant", "content": "2"},
            {"role": "user", "content": "3"},
            {"role": "assistant", "content": "4"},
            {"role": "user", "content": "5"},
        ]
        groups = summarizer._find_summarization_candidates(messages, protected=set())
        # Should have one group spanning all messages
        assert len(groups) == 1
        assert groups[0] == (0, 4)

    def test_protected_splits_groups(self) -> None:
        """Protected messages split the candidates into groups."""
        summarizer = ProgressiveSummarizer(min_messages_to_summarize=2)
        messages = [
            {"role": "user", "content": "1"},
            {"role": "assistant", "content": "2"},
            {"role": "user", "content": "3"},  # Protected at index 2
            {"role": "assistant", "content": "4"},
            {"role": "user", "content": "5"},
            {"role": "assistant", "content": "6"},
        ]
        groups = summarizer._find_summarization_candidates(messages, protected={2})
        # Should have two groups: (0,1) and (3,5)
        assert len(groups) == 2
        assert groups[0] == (0, 1)
        assert groups[1] == (3, 5)

    def test_min_messages_filter(self) -> None:
        """Groups smaller than min_messages_to_summarize are filtered."""
        summarizer = ProgressiveSummarizer(min_messages_to_summarize=3)
        messages = [
            {"role": "user", "content": "1"},
            {"role": "assistant", "content": "2"},
            {"role": "user", "content": "3"},  # Protected
            {"role": "assistant", "content": "4"},
        ]
        groups = summarizer._find_summarization_candidates(messages, protected={2})
        # Group (0,1) has 2 messages, filtered. Group (3,3) has 1, filtered.
        assert len(groups) == 0

    def test_all_protected_no_candidates(self) -> None:
        """No candidates when all messages are protected."""
        summarizer = ProgressiveSummarizer(min_messages_to_summarize=1)
        messages = [
            {"role": "user", "content": "1"},
            {"role": "assistant", "content": "2"},
        ]
        groups = summarizer._find_summarization_candidates(messages, protected={0, 1})
        assert len(groups) == 0

    def test_empty_messages(self) -> None:
        """Empty message list returns no groups."""
        summarizer = ProgressiveSummarizer()
        groups = summarizer._find_summarization_candidates([], protected=set())
        assert len(groups) == 0


# =============================================================================
# Summarize Messages Tests
# =============================================================================


class TestSummarizeMessages:
    """Tests for the main summarize_messages method."""

    def test_no_candidates_returns_original(
        self, tokenizer: Tokenizer, simple_conversation: list[dict[str, Any]]
    ) -> None:
        """When no candidates, return original messages unchanged."""
        summarizer = ProgressiveSummarizer(min_messages_to_summarize=100)  # Too high
        result = summarizer.summarize_messages(
            simple_conversation, tokenizer, protected_indices=set()
        )
        assert len(result.messages) == len(simple_conversation)
        assert result.tokens_saved == 0
        assert len(result.summaries_created) == 0

    def test_all_protected_no_changes(
        self, tokenizer: Tokenizer, simple_conversation: list[dict[str, Any]]
    ) -> None:
        """All protected messages means no summarization."""
        summarizer = ProgressiveSummarizer(min_messages_to_summarize=2)
        all_protected = set(range(len(simple_conversation)))
        result = summarizer.summarize_messages(
            simple_conversation, tokenizer, protected_indices=all_protected
        )
        assert len(result.messages) == len(simple_conversation)
        assert result.tokens_saved == 0

    def test_summarization_reduces_messages(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Summarization reduces message count."""
        summarizer = ProgressiveSummarizer(
            min_messages_to_summarize=3,
            store_for_retrieval=False,  # Skip CCR for test
        )
        # Protect first and last few messages
        protected = {0, 1, len(long_conversation) - 1, len(long_conversation) - 2}
        result = summarizer.summarize_messages(
            long_conversation, tokenizer, protected_indices=protected
        )

        # Should have fewer messages
        assert len(result.messages) < len(long_conversation)
        # Should save tokens
        assert result.tokens_saved > 0
        # Should create summaries
        assert len(result.summaries_created) > 0

    def test_summarization_result_structure(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Verify SummarizationResult has correct structure."""
        summarizer = ProgressiveSummarizer(
            min_messages_to_summarize=3,
            store_for_retrieval=False,
        )
        result = summarizer.summarize_messages(long_conversation, tokenizer, protected_indices={0})

        assert isinstance(result, SummarizationResult)
        assert isinstance(result.messages, list)
        assert isinstance(result.summaries_created, list)
        assert isinstance(result.tokens_before, int)
        assert isinstance(result.tokens_after, int)
        assert isinstance(result.transforms_applied, list)
        assert result.tokens_before >= result.tokens_after

    def test_custom_summarizer_called(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Custom summarizer function is called."""
        calls: list[int] = []

        def tracking_summarizer(messages: list[dict], context: str = "") -> str:
            calls.append(len(messages))
            return f"CUSTOM SUMMARY of {len(messages)} messages"

        summarizer = ProgressiveSummarizer(
            summarize_fn=tracking_summarizer,
            min_messages_to_summarize=3,
            store_for_retrieval=False,
        )

        result = summarizer.summarize_messages(long_conversation, tokenizer, protected_indices={0})

        # Custom summarizer should have been called
        assert len(calls) > 0
        # Summary should appear in messages
        found_custom = any("CUSTOM SUMMARY" in msg.get("content", "") for msg in result.messages)
        assert found_custom

    def test_context_passed_to_summarizer(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Context messages are passed to summarizer."""
        received_context: list[str] = []

        def context_tracking_summarizer(messages: list[dict], context: str = "") -> str:
            received_context.append(context)
            return "Summary"

        summarizer = ProgressiveSummarizer(
            summarize_fn=context_tracking_summarizer,
            min_messages_to_summarize=3,
            store_for_retrieval=False,
        )

        context_msgs = [{"role": "user", "content": "Recent important question"}]
        summarizer.summarize_messages(
            long_conversation,
            tokenizer,
            protected_indices={0},
            context_messages=context_msgs,
        )

        # Context should have been passed
        assert len(received_context) > 0
        # Should contain the recent message content
        assert any("Recent important question" in ctx for ctx in received_context)

    def test_target_tokens_stops_early(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Summarization stops when target tokens reached."""
        summarizer = ProgressiveSummarizer(
            min_messages_to_summarize=3,
            store_for_retrieval=False,
        )

        # Get original token count
        original_tokens = tokenizer.count_messages(long_conversation)

        # Set target very close to original (minimal summarization needed)
        target = int(original_tokens * 0.95)  # Only need 5% reduction

        result = summarizer.summarize_messages(
            long_conversation,
            tokenizer,
            protected_indices={0},
            target_tokens=target,
        )

        # Should stop once target reached
        assert result.tokens_after <= target or result.tokens_after < original_tokens

    def test_small_groups_skipped(
        self,
        tokenizer: Tokenizer,
    ) -> None:
        """Groups with < 100 tokens are skipped."""
        # Very short messages
        messages = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Hi"},
            {"role": "user", "content": "Bye"},
            {"role": "assistant", "content": "Bye"},
        ]

        summarizer = ProgressiveSummarizer(
            min_messages_to_summarize=2,
            store_for_retrieval=False,
        )

        result = summarizer.summarize_messages(messages, tokenizer, protected_indices=set())

        # Small groups should be skipped
        assert len(result.summaries_created) == 0

    def test_summary_larger_than_original_skipped(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Summaries larger than original are skipped."""

        def verbose_summarizer(messages: list[dict], context: str = "") -> str:
            # Return a very verbose summary
            return "VERY LONG SUMMARY " * 1000

        summarizer = ProgressiveSummarizer(
            summarize_fn=verbose_summarizer,
            min_messages_to_summarize=3,
            store_for_retrieval=False,
        )

        result = summarizer.summarize_messages(long_conversation, tokenizer, protected_indices={0})

        # Summaries larger than original should be skipped
        # (or if any were created, they saved tokens)
        for summary in result.summaries_created:
            assert summary.tokens_saved >= 0

    def test_summarizer_exception_handled(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Exceptions from summarizer are handled gracefully."""

        def failing_summarizer(messages: list[dict], context: str = "") -> str:
            raise ValueError("Summarization failed!")

        summarizer = ProgressiveSummarizer(
            summarize_fn=failing_summarizer,
            min_messages_to_summarize=3,
            store_for_retrieval=False,
        )

        # Should not raise, should return original
        result = summarizer.summarize_messages(long_conversation, tokenizer, protected_indices={0})

        # No summaries created due to failures
        assert len(result.summaries_created) == 0


# =============================================================================
# Integration Tests
# =============================================================================


class TestProgressiveSummarizerIntegration:
    """Integration tests for end-to-end summarization."""

    def test_full_workflow_with_extractive(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Test full workflow with default extractive summarizer."""
        summarizer = ProgressiveSummarizer(
            min_messages_to_summarize=4,
            store_for_retrieval=False,
        )

        original_count = len(long_conversation)

        result = summarizer.summarize_messages(
            long_conversation,
            tokenizer,
            protected_indices={0},  # Only protect system message
        )

        # Verify reduction
        assert len(result.messages) < original_count
        assert result.tokens_after < result.tokens_before

        # Verify transforms tracked
        assert len(result.transforms_applied) > 0

        # Verify summaries created
        assert len(result.summaries_created) > 0
        for summary in result.summaries_created:
            assert summary.start_index >= 0
            assert summary.end_index >= summary.start_index
            assert summary.compression_ratio < 1.0  # Actually compressed

    def test_preserves_protected_messages(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Protected messages are preserved exactly."""
        summarizer = ProgressiveSummarizer(
            min_messages_to_summarize=3,
            store_for_retrieval=False,
        )

        # Protect first 3 and last 3 messages
        protected = {
            0,
            1,
            2,
            len(long_conversation) - 3,
            len(long_conversation) - 2,
            len(long_conversation) - 1,
        }

        # Store original protected content
        original_protected = {i: long_conversation[i]["content"] for i in protected}

        result = summarizer.summarize_messages(
            long_conversation,
            tokenizer,
            protected_indices=protected,
        )

        # Find protected messages in result
        # First 3 should still be at beginning
        assert result.messages[0]["content"] == original_protected[0]
        assert result.messages[1]["content"] == original_protected[1]
        assert result.messages[2]["content"] == original_protected[2]

        # Last 3 should still be at end (positions shifted)
        assert result.messages[-1]["content"] == original_protected[len(long_conversation) - 1]
        assert result.messages[-2]["content"] == original_protected[len(long_conversation) - 2]
        assert result.messages[-3]["content"] == original_protected[len(long_conversation) - 3]

    def test_tool_messages_handled(
        self, tokenizer: Tokenizer, conversation_with_tools: list[dict[str, Any]]
    ) -> None:
        """Tool messages are handled in summarization."""
        # Create longer tool-heavy conversation
        long_tool_conv = conversation_with_tools.copy()
        for i in range(10):
            long_tool_conv.extend(
                [
                    {"role": "user", "content": f"Search again {i}"},
                    {
                        "role": "assistant",
                        "content": f"Searching {i}...",
                        "tool_calls": [
                            {
                                "id": f"call_{i}",
                                "type": "function",
                                "function": {"name": "search", "arguments": "{}"},
                            }
                        ],
                    },
                    {
                        "role": "tool",
                        "tool_call_id": f"call_{i}",
                        "content": f'{{"data": "result {i}"}}',
                    },
                    {"role": "assistant", "content": f"Found result {i}"},
                ]
            )

        summarizer = ProgressiveSummarizer(
            min_messages_to_summarize=3,
            store_for_retrieval=False,
        )

        result = summarizer.summarize_messages(
            long_tool_conv,
            tokenizer,
            protected_indices={0},
        )

        # Should reduce messages
        assert len(result.messages) < len(long_tool_conv)

        # Tool names should be tracked in summaries
        all_tool_names = []
        for summary in result.summaries_created:
            all_tool_names.extend(summary.tool_names)
        # Some tool calls should be tracked (may be empty if extractive)

    def test_does_not_mutate_original(
        self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
    ) -> None:
        """Original messages are not mutated."""
        import copy

        original_copy = copy.deepcopy(long_conversation)

        summarizer = ProgressiveSummarizer(
            min_messages_to_summarize=3,
            store_for_retrieval=False,
        )

        summarizer.summarize_messages(
            long_conversation,
            tokenizer,
            protected_indices={0},
        )

        # Original should be unchanged
        assert long_conversation == original_copy


# =============================================================================
# SummarizationResult Tests
# =============================================================================


class TestSummarizationResult:
    """Tests for SummarizationResult dataclass."""

    def test_tokens_saved_property(self) -> None:
        """Test tokens_saved property."""
        result = SummarizationResult(
            messages=[],
            summaries_created=[],
            tokens_before=1000,
            tokens_after=300,
            transforms_applied=[],
        )
        assert result.tokens_saved == 700

    def test_tokens_saved_no_negative(self) -> None:
        """Test tokens_saved doesn't go negative."""
        result = SummarizationResult(
            messages=[],
            summaries_created=[],
            tokens_before=100,
            tokens_after=150,
            transforms_applied=[],
        )
        assert result.tokens_saved == 0