"""Unit tests for context compressor."""

import pytest

from navi.core.compressor import (
    ContextCompressor,
    _build_summary_system_prompt,
    _format_for_summary,
    _turn_importance,
    compress_context,
    partition_messages,
    should_compress,
)
from navi.llm.base import Message, ToolCallRequest
from tests.conftest_factory import FakeLLMBackend, make_profile


class TestShouldCompress:
    def test_below_threshold(self):
        assert should_compress(100, 1000, 0.7) is False

    def test_at_threshold(self):
        assert should_compress(700, 1000, 0.7) is True

    def test_above_threshold(self):
        assert should_compress(800, 1000, 0.7) is True


class TestPartitionMessages:
    def test_empty(self):
        old, recent = partition_messages([], keep_recent=2)
        assert old == []
        assert recent == []

    def test_fewer_turns_than_keep(self):
        msgs = [
            Message(role="user", content="hi"),
            Message(role="assistant", content="hello"),
        ]
        old, recent = partition_messages(msgs, keep_recent=5)
        assert old == []
        assert recent == msgs

    def test_splits_into_old_and_recent(self):
        msgs = [
            Message(role="user", content="1"),
            Message(role="assistant", content="a1"),
            Message(role="user", content="2"),
            Message(role="assistant", content="a2"),
            Message(role="user", content="3"),
            Message(role="assistant", content="a3"),
        ]
        old, recent = partition_messages(msgs, keep_recent=2)
        assert len(old) == 2  # turn 1
        assert len(recent) == 4  # turns 2+3

    def test_system_messages_ignored(self):
        msgs = [
            Message(role="system", content="sys"),
            Message(role="user", content="1"),
            Message(role="assistant", content="a1"),
        ]
        old, recent = partition_messages(msgs, keep_recent=5)
        assert recent == [m for m in msgs if m.role != "system"]

    def test_tool_calls_stay_with_assistant(self):
        msgs = [
            Message(role="user", content="1"),
            Message(
                role="assistant",
                content="",
                tool_calls=[ToolCallRequest(id="1", name="fs", arguments={})],
            ),
            Message(role="tool", content="result", name="fs", tool_call_id="1"),
            Message(role="user", content="2"),
            Message(role="assistant", content="ok"),
        ]
        old, recent = partition_messages(msgs, keep_recent=1)
        assert len(old) == 3
        assert len(recent) == 2

    def test_mid_turn_fallback_keeps_user_and_recent_messages(self):
        msgs = [Message(role="user", content="build a model")]
        for i in range(5):
            msgs.append(Message(
                role="assistant",
                tool_calls=[ToolCallRequest(id=str(i), name="fs", arguments={})],
            ))
            msgs.append(Message(role="tool", content=f"result {i}", name="fs", tool_call_id=str(i)))

        old, recent = partition_messages(msgs, keep_recent=8, keep_recent_messages=4)

        assert old
        assert recent[0].role == "user"
        assert recent[0].content == "build a model"
        assert len(recent) == 5  # original user message + 4 newest in-flight messages

    def test_midturn_splits_in_flight_turn_when_many_turns(self):
        """When turns > keep_recent AND keep_recent_messages is set, the in-flight
        (last) turn is also split — head (user request) + newest K messages kept,
        middle summarized — on top of compressing older turns. Previously the
        in-flight turn stayed verbatim in the keep window (the 150->141 shallow
        midturn-compression symptom)."""
        msgs = []
        # 14 older short turns (2 msgs each) = 28 msgs
        for i in range(14):
            msgs.append(Message(role="user", content=f"u{i}"))
            msgs.append(Message(role="assistant", content=f"a{i}"))
        # in-flight 15th turn: user + 10 assistant/tool pairs = 21 msgs
        msgs.append(Message(role="user", content="inflight request"))
        for i in range(10):
            msgs.append(Message(role="assistant", content=f"step {i}"))
            msgs.append(Message(role="tool", content=f"res {i}", name="fs", tool_call_id=str(i)))

        old, recent = partition_messages(msgs, keep_recent=12, keep_recent_messages=4)

        recent_contents = {m.content for m in recent}
        old_contents = {m.content for m in old}

        # The in-flight user request is kept verbatim (head of the in-flight turn).
        assert "inflight request" in recent_contents
        # The newest in-flight messages are kept verbatim (tail).
        assert "step 9" in recent_contents and "res 9" in recent_contents
        # The oldest older turn is summarized.
        assert "u0" in old_contents and "a0" in old_contents
        # The middle of the in-flight turn is summarized, not kept verbatim.
        assert "step 0" in old_contents and "res 0" in old_contents
        assert "step 0" not in recent_contents
        # Nothing is dropped: old + recent partition the whole non-system set.
        assert len(old) + len(recent) == len(msgs)


class TestFormatForSummary:
    def test_user_message(self):
        msgs = [Message(role="user", content="hello")]
        text, images = _format_for_summary(msgs)
        assert "User: hello" in text
        assert images == []

    def test_assistant_message(self):
        msgs = [Message(role="assistant", content="world")]
        text, images = _format_for_summary(msgs)
        assert "Assistant: world" in text

    def test_tool_call_block(self):
        msgs = [
            Message(
                role="assistant",
                tool_calls=[ToolCallRequest(id="1", name="fs", arguments={"path": "/tmp"})],
            ),
            Message(role="tool", content="file data", name="fs", tool_call_id="1"),
        ]
        text, images = _format_for_summary(msgs)
        assert "[Tool call: fs" in text
        assert "[Tool result: fs" in text

    def test_images_collected(self):
        msgs = [Message(role="user", content="look", images=["base64img"])]
        text, images = _format_for_summary(msgs)
        assert images == ["base64img"]
        assert "[+ 1 image(s)]" in text

    def test_summary_message_folded(self):
        msgs = [Message(role="user", content="old summary", is_summary=True)]
        text, _ = _format_for_summary(msgs)
        assert "old summary" in text


class TestCompressContext:
    async def test_nothing_to_compress(self):
        backend = FakeLLMBackend()
        result = await compress_context(
            context=[Message(role="user", content="hi")],
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=8,
        )
        assert result is None

    async def test_compresses_old_turns(self):
        backend = FakeLLMBackend(responses=["Summary of old stuff"])
        context = [
            Message(role="system", content="sys"),
            Message(role="user", content="1"),
            Message(role="assistant", content="a1"),
            Message(role="user", content="2"),
            Message(role="assistant", content="a2"),
            Message(role="user", content="3"),
            Message(role="assistant", content="a3"),
        ]
        new_context, summary = await compress_context(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=2,
        )
        assert summary == "Summary of old stuff"
        # system + summary + 2 recent turns (user+assistant × 2) = 6
        assert len(new_context) == 6
        assert new_context[0].role == "system"
        assert new_context[1].is_summary is True
        assert new_context[1].is_display is False
        assert new_context[1].is_context is True

    async def test_preserves_system_messages(self):
        backend = FakeLLMBackend(responses=["sum"])
        context = [
            Message(role="system", content="s1"),
            Message(role="system", content="s2"),
            Message(role="user", content="1"),
            Message(role="assistant", content="a1"),
            Message(role="user", content="2"),
            Message(role="assistant", content="a2"),
        ]
        new_context, _ = await compress_context(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=1,
        )
        system_msgs = [m for m in new_context if m.role == "system"]
        assert len(system_msgs) == 2

    async def test_compresses_long_current_turn_when_requested(self):
        backend = FakeLLMBackend(responses=["mid-turn summary"])
        context = [Message(role="user", content="build a model")]
        for i in range(5):
            context.append(Message(
                role="assistant",
                tool_calls=[ToolCallRequest(id=str(i), name="fs", arguments={})],
            ))
            context.append(Message(role="tool", content=f"large result {i}", name="fs", tool_call_id=str(i)))

        new_context, summary = await compress_context(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=8,
            keep_recent_messages=4,
        )

        assert summary == "mid-turn summary"
        assert new_context[0].is_summary is True
        assert new_context[1].role == "user"
        assert new_context[1].content == "build a model"
        assert len(new_context) == 6  # summary + user + 4 recent messages

    async def test_meta_summary_consolidates_multiple_summaries(self):
        """When to_summarize contains multiple long existing summaries, a meta-summary
        pass runs first (consolidating them) before the main compression."""
        # First response = meta-summary, second = main compression
        backend = FakeLLMBackend(responses=["Meta summary", "Final summary"])
        # Build a context with two existing summaries (each > 4000 chars to cross threshold)
        big_summary_1 = "A" * 5000
        big_summary_2 = "B" * 5000
        context = [
            Message(role="system", content="sys"),
            Message(role="user", content=big_summary_1, is_summary=True, is_display=False),
            Message(role="user", content=big_summary_2, is_summary=True, is_display=False),
            Message(role="user", content="recent question"),
            Message(role="assistant", content="recent answer"),
            Message(role="user", content="new question"),
            Message(role="assistant", content="new answer"),
        ]
        new_context, summary = await compress_context(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=2,
        )
        assert summary == "Final summary"
        # system + 1 consolidated summary + 2 recent turns = 6
        assert len(new_context) == 6
        assert new_context[1].is_summary is True
        # Exactly one summary in final context (meta + raw folded into one)
        assert sum(1 for m in new_context if m.is_summary) == 1
        # Two LLM calls happened (meta + main)
        assert backend._call_idx == 2

    async def test_meta_summary_skipped_when_summaries_are_short(self):
        """Short existing summaries should not trigger an extra meta-summary pass."""
        backend = FakeLLMBackend(responses=["Final summary"])
        context = [
            Message(role="system", content="sys"),
            Message(role="user", content="short summary 1", is_summary=True, is_display=False),
            Message(role="user", content="short summary 2", is_summary=True, is_display=False),
            Message(role="user", content="recent question"),
            Message(role="assistant", content="recent answer"),
            Message(role="user", content="new question"),
            Message(role="assistant", content="new answer"),
        ]
        new_context, summary = await compress_context(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=2,
        )
        assert summary == "Final summary"
        assert len(new_context) == 6
        # Only one LLM call (main compression) because summaries were too short
        # to trigger meta-summary
        assert backend._call_idx == 1

    async def test_meta_summary_graceful_on_failure(self):
        """If meta-summary fails, compression continues with raw summaries."""
        import navi.core.compressor as compressor_module

        backend = FakeLLMBackend(responses=["Final summary"])
        original_meta = compressor_module._meta_summarize

        async def _failing_meta(*args, **kwargs):
            raise RuntimeError("meta boom")

        compressor_module._meta_summarize = _failing_meta
        try:
            big = "X" * 5000
            context = [
                Message(role="system", content="sys"),
                Message(role="user", content=big, is_summary=True, is_display=False),
                Message(role="user", content=big, is_summary=True, is_display=False),
                Message(role="user", content="q1"),
                Message(role="assistant", content="a1"),
                Message(role="user", content="q2"),
                Message(role="assistant", content="a2"),
            ]
            new_context, summary = await compress_context(
                context=context,
                llm=backend,
                model="test",
                temperature=0.3,
                keep_recent=2,
            )
            assert summary == "Final summary"
            assert len(new_context) == 6
            # Meta failed, so only one call to the backend
            assert backend._call_idx == 1
        finally:
            compressor_module._meta_summarize = original_meta

    async def test_intra_turn_fallback_aggressive(self):
        """When turn-based partition has nothing to compress but keep_recent_messages
        is set, an aggressive fallback (keep_recent_messages=2) should still find
        something to summarize.

        Here the current turn is exactly 5 messages (user + 2 assistant/tool pairs).
        With keep_recent_messages=4, partition_current_turn_messages returns None
        because len(turn) <= keep_recent_messages + 1. This forces the fallback to
        retry with keep_recent_messages=2, which successfully compresses."""
        backend = FakeLLMBackend(responses=["aggressive summary"])
        context = [Message(role="user", content="autonomous loop")]
        for i in range(2):
            context.append(Message(
                role="assistant",
                tool_calls=[ToolCallRequest(id=str(i), name="fs", arguments={})],
            ))
            context.append(Message(role="tool", content=f"result {i}", name="fs", tool_call_id=str(i)))

        # Single turn, keep_recent=1 → turn-based partition returns ([], non_system)
        # Fallback should kick in with keep_recent_messages=2
        new_context, summary = await compress_context(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=1,
            keep_recent_messages=4,
        )

        assert summary == "aggressive summary"
        assert new_context[0].is_summary is True
        assert new_context[1].role == "user"
        # Fallback with keep_recent_messages=2 keeps user + 2 newest messages
        assert len(new_context) == 4  # summary + user + 2 recent messages


class TestContextCompressor:
    def test_estimate_context_tokens_text_only(self):
        compressor = ContextCompressor()
        context = [Message(role="user", content="hello world")]
        assert compressor.estimate_context_tokens(context) == 3  # 11 chars // 3 = 3

    def test_estimate_context_tokens_with_images(self):
        compressor = ContextCompressor()
        context = [Message(role="user", content="hi", images=["base64img"])]
        assert compressor.estimate_context_tokens(context) == 500  # 2 chars // 3 = 0 + 500

    @pytest.mark.asyncio
    async def test_compress_session_success(self):
        backend = FakeLLMBackend(responses=["Summary text"])
        compressor = ContextCompressor()
        context = [
            Message(role="system", content="sys"),
            Message(role="user", content="1"),
            Message(role="assistant", content="a1"),
            Message(role="user", content="2"),
            Message(role="assistant", content="a2"),
            Message(role="user", content="3"),
            Message(role="assistant", content="a3"),
        ]
        result = await compressor.compress_session(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=2,
        )
        assert result is not None
        new_context, summary = result
        assert summary == "Summary text"
        assert len(new_context) == 6  # system + summary + 2 recent turns

    @pytest.mark.asyncio
    async def test_midturn_compresses_in_flight_turn_with_many_turns(self):
        """End-to-end midturn: with > keep_recent turns and a long in-flight turn,
        compression shrinks the in-flight turn too (head + newest K verbatim, middle
        summarized) — not just the older turns. Previously the in-flight turn stayed
        verbatim and compression was shallow."""
        backend = FakeLLMBackend(responses=["consolidated summary"])
        compressor = ContextCompressor()
        context = []
        for i in range(14):
            context.append(Message(role="user", content=f"u{i}"))
            context.append(Message(role="assistant", content=f"a{i}"))
        context.append(Message(role="user", content="inflight request"))
        for i in range(10):
            context.append(Message(role="assistant", content=f"step {i}"))
            context.append(Message(role="tool", content=f"res {i}", name="fs", tool_call_id=str(i)))

        result = await compressor.compress_session(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=12,
            keep_recent_messages=4,
        )
        assert result is not None
        new_context, summary = result
        assert summary == "consolidated summary"
        # 1 summary + head(user) + 11 kept older turns * 2 msgs + 4 in-flight tail
        # = 1 + 1 + 22 + 4 = 28. Without A1 the in-flight turn (21 msgs) would stay
        # verbatim and new_context would be ~1 + 24 + 21 = 46.
        assert len(new_context) == 28
        assert new_context[0].is_summary is True
        contents = {m.content for m in new_context}
        assert "inflight request" in contents  # head kept
        assert "step 9" in contents and "res 9" in contents  # tail kept
        assert "step 0" not in contents  # in-flight middle summarized away

    @pytest.mark.asyncio
    async def test_compress_session_retry_on_failure(self):
        """First attempt fails, second succeeds with keep_recent + 4."""
        import navi.core.compressor as compressor_module

        backend = FakeLLMBackend(responses=["Summary text"])
        original_compress_context = compressor_module.compress_context
        call_count = 0

        async def _failing_once(*args, **kwargs):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                raise RuntimeError("boom")
            return await original_compress_context(*args, **kwargs)

        compressor_module.compress_context = _failing_once
        try:
            compressor = ContextCompressor()
            context = [Message(role="system", content="sys")]
            for i in range(7):
                context.append(Message(role="user", content=str(i)))
                context.append(Message(role="assistant", content=f"a{i}"))
            result = await compressor.compress_session(
                context=context,
                llm=backend,
                model="test",
                temperature=0.3,
                keep_recent=2,
            )
            assert result is not None
            assert call_count == 2
        finally:
            compressor_module.compress_context = original_compress_context

    @pytest.mark.asyncio
    async def test_compress_session_hard_truncate_on_double_failure(self):
        """Both attempts fail → token-budget hard-truncate drops oldest turns."""
        import navi.core.compressor as compressor_module
        from unittest.mock import AsyncMock

        original_compress_context = compressor_module.compress_context

        async def _always_fail(*args, **kwargs):
            raise RuntimeError("always fails")

        compressor_module.compress_context = _always_fail
        try:
            compressor = ContextCompressor()
            # Four turns, each large enough that the whole context far exceeds
            # the hard-truncate token budget (num_ctx * 0.5). The fallback must
            # keep only the newest turn(s) that fit the budget and drop the rest.
            big = "x" * 120_000  # ~40k tokens each — two already overflow 0.5*65k
            context = [
                Message(role="system", content="sys"),
                Message(role="user", content="1"), Message(role="assistant", content=big),
                Message(role="user", content="2"), Message(role="assistant", content=big),
                Message(role="user", content="3"), Message(role="assistant", content=big),
                Message(role="user", content="4"), Message(role="assistant", content=big),
            ]
            result = await compressor.compress_session(
                context=context,
                llm=AsyncMock(),
                model="test",
                temperature=0.3,
                keep_recent=2,
            )
            assert result is not None
            new_context, summary = result
            assert "truncated" in summary.lower()
            # System + only the newest turn (user "4" + its assistant) fit the
            # budget; the three older turns are dropped.
            assert len(new_context) == 3
            assert new_context[0].role == "system"
            assert new_context[1].content == "4"
        finally:
            compressor_module.compress_context = original_compress_context

    @pytest.mark.asyncio
    async def test_compress_session_returns_none_when_nothing_to_compress(self):
        backend = FakeLLMBackend()
        compressor = ContextCompressor()
        context = [
            Message(role="user", content="hi"),
            Message(role="assistant", content="hello"),
        ]
        result = await compressor.compress_session(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=5,
        )
        assert result is None

    def test_turn_importance(self):
        important = [
            Message(role="user", content="that is wrong, use json instead"),
        ]
        casual = [
            Message(role="user", content="hi"),
            Message(role="assistant", content="hello"),
        ]
        assert _turn_importance(important) > _turn_importance(casual)

    def test_adaptive_partition_keeps_important_turns(self):
        msgs = []
        for i in range(6):
            msgs.append(Message(role="user", content=f"task {i}"))
            msgs.append(Message(role="assistant", content=f"answer {i}"))
        # Mark one old turn as important; it should survive even with keep_recent=2
        msgs[6].content = "this is wrong, fix it"
        old, recent = partition_messages(msgs, keep_recent=2)
        assert msgs[6] in recent  # user correction stays

    def test_format_keeps_critical_tool_result(self):
        long_result = "exact error output\n" * 100
        msgs = [
            Message(
                role="assistant",
                tool_calls=[ToolCallRequest(id="1", name="filesystem", arguments={})],
            ),
            Message(role="tool", content=long_result, name="filesystem", tool_call_id="1", is_compression_critical=True),
        ]
        text, _ = _format_for_summary(msgs)
        assert long_result in text

    def test_format_truncates_noncritical_tool_result(self):
        long_result = "x" * 1000
        msgs = [
            Message(
                role="assistant",
                tool_calls=[ToolCallRequest(id="1", name="web_search", arguments={})],
            ),
            Message(role="tool", content=long_result, name="web_search", tool_call_id="1"),
        ]
        text, _ = _format_for_summary(msgs)
        assert long_result not in text
        assert "…" in text

    def test_summary_prompt_uses_profile_compression_prompt_file(self):
        profile = make_profile("navi_code", compression_prompt_file="compression_prompt.txt")
        prompt = _build_summary_system_prompt(profile)
        assert "## Goal" in prompt
        assert "## Active Files" in prompt

    async def test_profile_overrides_compression_max_tokens(self):
        profile = make_profile("test", compression_max_tokens=1234)
        backend = FakeLLMBackend(responses=["short"])
        context = [
            Message(role="user", content="1"),
            Message(role="assistant", content="a1"),
            Message(role="user", content="2"),
            Message(role="assistant", content="a2"),
            Message(role="user", content="3"),
            Message(role="assistant", content="a3"),
        ]
        _, _ = await compress_context(
            context=context,
            llm=backend,
            model="test",
            temperature=0.3,
            keep_recent=1,
            max_tokens=9999,
            profile=profile,
        )
        # max_tokens passed to llm.complete should be overridden by profile
        # FakeLLMBackend ignores max_tokens, but we can at least verify the call ran
        assert backend._call_idx == 1


class TestTargetHysteresis:
    """Target hysteresis: after compression the kept region must fit the target
    fraction (~65%) of the context window, not just skim below the 90% trigger —
    otherwise the next few messages re-fire compression and the context yo-yos
    at the trigger line ("125 -> 104, then 104 again after a couple of messages").
    """

    @staticmethod
    def _override_settings(ollama_num_ctx: int):
        """Patch both navi.config.settings and navi.core.compressor.settings.

        compress_context reads the module-level ``settings`` binding of
        compressor.py; compress_session's safety net reads ``navi.config.settings``
        via a local ``import navi.config as _cfg``. Both must be patched for the
        override to reach every code path.
        """
        import navi.config as _config
        import navi.core.compressor as _cm
        from navi.config import Settings

        new = Settings(
            _env_file=None,
            ollama_num_ctx=ollama_num_ctx,
            output_reserve_tokens=2048,
            navi_persona_file="",
            context_compression_target=0.65,
        )
        prev_config = _config.settings
        prev_cm = _cm.settings
        _config.settings = new
        _cm.settings = new
        return prev_config, prev_cm

    @staticmethod
    def _restore(prev_config, prev_cm):
        import navi.config as _config
        import navi.core.compressor as _cm

        _config.settings = prev_config
        _cm.settings = prev_cm

    @pytest.mark.asyncio
    async def test_preturn_adaptive_shrink_fits_target(self):
        """Many large turns, turn-based (preturn) mode: a fixed keep_recent would
        leave the kept region above the 65% target, so compress_context shrinks
        keep_recent until the verbatim kept region fits. The extra turns fold
        into the summary (still one LLM call)."""
        prev_config, prev_cm = self._override_settings(ollama_num_ctx=10_000)
        try:
            # target_tokens = 10_000 * 0.65 = 6_500
            backend = FakeLLMBackend(responses=["Summary text"])
            context = [Message(role="system", content="sys")]
            # 20 turns, each assistant ~2500 chars (~833 tokens) -> a turn ~835
            # tokens. 12 turns ~10_020 > 6_500; 7 turns ~5_845 < 6_500.
            for i in range(20):
                context.append(Message(role="user", content=f"u{i}"))
                context.append(Message(role="assistant", content="a" * 2500))

            new_context, _ = await compress_context(
                context=context,
                llm=backend,
                model="test",
                temperature=0.3,
                keep_recent=12,
            )
            assert new_context is not None
            # system + summary + 7 kept turns * 2 msgs = 16 (not 1+1+24=26)
            assert len(new_context) == 16
            # The kept region fits the target.
            assert ContextCompressor.estimate_context_tokens(new_context) <= 6_500
            # Newest turns kept verbatim, oldest summarized away.
            contents = {m.content for m in new_context}
            assert "u19" in contents and "u13" in contents  # last 7 turns kept
            assert "u0" not in contents and "u12" not in contents  # summarized
            # Single LLM call — shrink used partition only, no extra summarization.
            assert backend._call_idx == 1
        finally:
            self._restore(prev_config, prev_cm)

    @pytest.mark.asyncio
    async def test_preturn_no_shrink_when_already_under_target(self):
        """Small context: the kept region already fits the target, so
        keep_recent is not shrunk — normal compression behavior is preserved."""
        prev_config, prev_cm = self._override_settings(ollama_num_ctx=100_000)
        try:
            backend = FakeLLMBackend(responses=["Summary text"])
            context = [
                Message(role="system", content="sys"),
                Message(role="user", content="1"),
                Message(role="assistant", content="a1"),
                Message(role="user", content="2"),
                Message(role="assistant", content="a2"),
                Message(role="user", content="3"),
                Message(role="assistant", content="a3"),
            ]
            new_context, _ = await compress_context(
                context=context,
                llm=backend,
                model="test",
                temperature=0.3,
                keep_recent=2,
            )
            # system + summary + 2 kept turns = 6 — unchanged from before hysteresis.
            assert len(new_context) == 6
            assert ContextCompressor.estimate_context_tokens(new_context) <= 65_000
        finally:
            self._restore(prev_config, prev_cm)

    @pytest.mark.asyncio
    async def test_safety_net_truncates_when_floor_still_over_target(self):
        """Even the keep_recent floor (2 turns) exceeds the target — the verbatim
        kept region cannot fit. compress_context returns the floor result (over
        target), and compress_session's safety net then drops oldest kept turns
        by token budget so a compression firing always leaves real headroom."""
        prev_config, prev_cm = self._override_settings(ollama_num_ctx=10_000)
        try:
            # target_tokens = 6_500. 6 turns, each assistant ~12_000 chars
            # (~4_000 tokens) -> 2 turns ~8_000 > 6_500, so the floor (2) is
            # still over target. compress_session must truncate further.
            big = "x" * 12_000
            backend = FakeLLMBackend(responses=["Summary text"])
            compressor = ContextCompressor()
            context = [Message(role="system", content="sys")]
            for i in range(6):
                context.append(Message(role="user", content=f"u{i}"))
                context.append(Message(role="assistant", content=big))

            result = await compressor.compress_session(
                context=context,
                llm=backend,
                model="test",
                temperature=0.3,
                keep_recent=4,
            )
            assert result is not None
            new_context, _ = result
            # Safety net brought the context under the target.
            assert compressor.estimate_context_tokens(new_context) <= 6_500
            # Only the newest turn fit the budget; older kept turns + the summary
            # were dropped (they were verbatim, not summarized — the documented
            # safety-net loss for the pathological "kept region alone exceeds
            # the target" case).
            assert len(new_context) == 3  # system + newest user + its assistant
            assert new_context[1].content == "u5"
            assert not any(m.is_summary for m in new_context)
        finally:
            self._restore(prev_config, prev_cm)


class TestWouldCompress:
    """would_compress must predict compress_session's actual outcome (no LLM call),
    so the agent only announces compression when it will really happen."""

    def test_false_for_small_context(self):
        c = ContextCompressor()
        ctx = [Message(role="user", content="hi"), Message(role="assistant", content="hello")]
        assert c.would_compress(ctx, keep_recent=8, threshold=0.70, max_context_tokens=65536) is False

    def test_true_when_partition_finds_messages(self):
        c = ContextCompressor()
        ctx = [Message(role="user", content="do task")] + [
            Message(role="assistant", content=f"s{i}") for i in range(20)
        ]
        # many messages in one turn -> intra-turn split finds to_summarize
        assert c.would_compress(
            ctx, keep_recent=8, keep_recent_messages=16, threshold=0.70, max_context_tokens=65536
        ) is True

    def test_false_for_single_huge_turn_that_cannot_be_split(self):
        """A few very large messages in one turn: the token gate would fire, but
        neither partition nor the token-budget truncate can shrink a single turn
        that alone exceeds the budget — so would_compress is False (the
        per-message view truncation in context_builder.build handles this case)."""
        c = ContextCompressor()
        big = "x" * 200_000
        ctx = [
            Message(role="user", content="read file"),
            Message(role="assistant", content="a"),
            Message(role="tool", name="fs", content=big),
        ]
        assert c.would_compress(
            ctx, keep_recent=12, keep_recent_messages=16, threshold=0.70, max_context_tokens=65536
        ) is False

    def test_true_when_token_budget_truncate_can_drop_oldest_turns(self):
        """Multiple huge turns: dropping oldest turns brings the context under
        budget, so the token-budget fallback will shrink it -> would_compress True."""
        c = ContextCompressor()
        big = "x" * 120_000
        ctx = []
        for i in range(4):
            ctx.append(Message(role="user", content=str(i)))
            ctx.append(Message(role="assistant", content=big))
        assert c.would_compress(
            ctx, keep_recent=12, keep_recent_messages=16, threshold=0.70, max_context_tokens=65536
        ) is True


class TestTokenBudgetFallback:
    """compress_session shrinks via token-budget truncate when partition no-ops
    but the context is over the threshold (many large turns)."""

    @pytest.mark.asyncio
    async def test_compress_session_drops_oldest_turns_via_token_budget(self):
        c = ContextCompressor()
        big = "x" * 120_000  # ~40k tokens each
        ctx = [Message(role="system", content="sys")]
        for i in range(4):
            ctx.append(Message(role="user", content=str(i)))
            ctx.append(Message(role="assistant", content=big))
        # llm=None is fine: partition no-ops (1 turn? no — 4 turns, but keep_recent=12
        # > 4 turns -> turn-based no-op; intra-turn krm=16 on the LAST turn only has
        # 2 msgs -> no-op too). So the token-budget fallback path is taken.
        result = await c.compress_session(
            context=ctx, llm=None, model="m", temperature=0.3,
            keep_recent=12, keep_recent_messages=16,
        )
        assert result is not None
        new_context, summary = result
        # system + only the newest turn fit the budget
        assert new_context[0].role == "system"
        assert len(new_context) == 3
        assert new_context[1].content == "3"

    @pytest.mark.asyncio
    async def test_compress_session_returns_none_when_view_truncation_suffices(self):
        """Single huge turn: token-budget truncate can't split it -> None.
        (context_builder.build truncates the per-message view instead.)"""
        c = ContextCompressor()
        big = "x" * 200_000
        ctx = [
            Message(role="user", content="read file"),
            Message(role="assistant", content="a"),
            Message(role="tool", name="fs", content=big),
        ]
        result = await c.compress_session(
            context=ctx, llm=None, model="m", temperature=0.3,
            keep_recent=12, keep_recent_messages=16,
        )
        assert result is None


class TestRealBaselineEstimate:
    """real_baseline_estimate uses real prompt_tokens for the bulk of the
    context + a heuristic delta for messages appended since, instead of the
    chars//3 estimate that undercounts code-heavy tool output."""

    def test_no_baseline_falls_back_to_heuristic(self):
        c = ContextCompressor()
        ctx = [Message(role="user", content="hello world")]  # 11 chars -> 3 tokens
        assert c.real_baseline_estimate(ctx, ctx) == 3

    def test_uses_real_baseline_plus_delta_for_new_messages(self):
        c = ContextCompressor()
        # Simulate: last LLM call saw 4 session messages and reported 5000 real
        # prompt tokens. Since then 2 new messages (12 chars -> 4 tokens) were
        # appended. Estimate should be 5000 + 4, NOT the chars//3 of the whole.
        base = [
            Message(role="user", content="do the task"),
            Message(role="assistant", content="step"),
            Message(role="tool", name="fs", content="result"),
            Message(role="assistant", content="ok"),
        ]
        c.record_real_baseline(len(base), 5000)
        session_context = base + [
            Message(role="tool", name="fs", content="newresult"),
            Message(role="assistant", content="done"),
        ]
        built = [Message(role="system", content="sys")] + session_context
        assert c.real_baseline_estimate(session_context, built) == 5004

    def test_stale_baseline_after_shrink_falls_back_to_heuristic(self):
        """After compression the context shrank below baseline_len -> the
        baseline is stale and must fall back to heuristic (not crash / neg)."""
        c = ContextCompressor()
        c.record_real_baseline(10, 5000)  # baseline recorded at len 10
        ctx = [Message(role="user", content="hi")]  # now len 1 < 10
        # falls back to heuristic on built_ctx
        assert c.real_baseline_estimate(ctx, ctx) == 0  # 2 chars // 3 = 0

    def test_record_real_baseline_ignores_none_and_zero(self):
        c = ContextCompressor()
        c.record_real_baseline(4, None)
        assert c._real_baseline is None
        c.record_real_baseline(4, 0)
        assert c._real_baseline is None
        c.record_real_baseline(4, 5000)
        assert c._real_baseline == (4, 5000)

    def test_check_context_size_uses_baseline_when_session_context_given(self):
        """check_context_size prefers the real baseline when the caller passes
        session_context; without it (subagent / unit tests) it stays heuristic."""
        import navi.config as _config
        from navi.config import Settings

        prev_settings = _config.settings
        _config.settings = Settings(
            _env_file=None, ollama_num_ctx=100_000, output_reserve_tokens=2048,
            navi_persona_file="",
        )
        try:
            c = ContextCompressor()
            # Real baseline says we are at 98k tokens; heuristic (chars//3) on a
            # tiny string would say ~0. With the baseline the guard must trip
            # (98k > 100k - 2048 reserve = 97952 available).
            base = [Message(role="user", content="task")]
            c.record_real_baseline(len(base), 98_000)
            session_context = base + [Message(role="assistant", content="ok")]
            built = [Message(role="system", content="sys")] + session_context
            from navi.exceptions import ContextTooLargeError

            with pytest.raises(ContextTooLargeError):
                c.check_context_size(built, session_context=session_context)
        finally:
            _config.settings = prev_settings
