"""Unit tests for navi.core.agent.Agent.

Uses InMemorySessionStore, FakeLLMBackend, and FakeTool so tests run
without a real database or LLM server.
"""

import asyncio
import copy
from datetime import datetime, timezone

import pytest
import pytest_asyncio

from navi.core.agent import Agent
from navi.core.events import (
    CompressionStarted,
    ContextCompressed,
    ModelInfo,
    StreamEnd,
    SubagentComplete,
)
from navi.core.registry import BackendRegistry, ProfileRegistry, ToolRegistry
from navi.core.session import InMemorySessionStore, Session
from navi.exceptions import MaxIterationsReached, NothingToCompactError, SessionNotFound
from navi.llm.base import LLMChunk, Message, ToolCallRequest
from tests.conftest_factory import (
    FakeLLMBackend,
    FakeTool,
    make_profile,
    make_registry_with_tools,
)


@pytest.fixture
def agent():
    sessions = InMemorySessionStore()
    profiles = ProfileRegistry()
    profile = make_profile("test")
    profiles.register(profile)
    tools = make_registry_with_tools()
    backends = BackendRegistry()
    backends.register("ollama", FakeLLMBackend(responses=["hello"]))
    return Agent(
        session_store=sessions,
        profile_registry=profiles,
        tool_registry=tools,
        backend_registry=backends,
    )


@pytest_asyncio.fixture
async def session(agent):
    return await agent._sessions.create(profile_id="test")


# ─── run() tests ───────────────────────────────────────────────────────────


class TestAgentRun:
    @pytest.mark.asyncio
    async def test_run_single_iteration(self, agent, session):
        backend = FakeLLMBackend(responses=["hello"])
        agent._backends.register("ollama", backend)

        result = await agent.run(session.id, "hi")
        assert result == "hello"
        saved = await agent._sessions.get(session.id)
        # user display + user context + assistant
        assert len(saved.messages) == 3
        assert saved.messages[0].role == "user"
        assert saved.messages[1].role == "user"
        assert saved.messages[2].role == "assistant"
        assert saved.messages[2].content == "hello"

        # Flags: display-only user message, context-only user message
        assert saved.messages[0].is_display is True
        assert saved.messages[0].is_context is False
        assert saved.messages[1].is_display is False
        assert saved.messages[1].is_context is True
        # Assistant message is both display and context
        assert saved.messages[2].is_display is True
        assert saved.messages[2].is_context is True

    @pytest.mark.asyncio
    async def test_run_session_not_found(self, agent):
        with pytest.raises(SessionNotFound):
            await agent.run("nonexistent-id", "hi")

    @pytest.mark.asyncio
    async def test_run_tool_calls_then_stop(self, agent, session):
        """Tool-calling turn followed by a final stop turn."""
        backend = FakeLLMBackend(
            responses=["", "done"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
        )
        agent._backends.register("ollama", backend)

        result = await agent.run(session.id, "do something")
        assert result == "done"
        saved = await agent._sessions.get(session.id)
        # user display + user context + assistant(tool) + tool_result + assistant(final)
        assert len(saved.messages) == 5
        assert saved.messages[3].role == "tool"
        assert saved.messages[4].content == "done"

    @pytest.mark.asyncio
    async def test_run_token_accumulation(self, agent, session):
        """_turn_tokens accumulates completion tokens across tool-calling iterations."""
        backend = FakeLLMBackend(
            responses=["", "done"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
            prompt_tokens=10,
            completion_tokens=5,
        )
        agent._backends.register("ollama", backend)

        await agent.run(session.id, "do something")
        saved = await agent._sessions.get(session.id)
        final_msg = saved.messages[-1]
        # Two iterations × 5 completion tokens = 10 tokens
        assert final_msg.token_count == 10

    @pytest.mark.asyncio
    async def test_run_max_iterations(self, agent, session):
        """After max_iterations tool turns, MaxIterationsReached is raised."""
        profile = agent._profiles.get("test")
        profile.max_iterations = 2

        backend = FakeLLMBackend(
            responses=["", ""],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                [ToolCallRequest(id="2", name="test_tool", arguments={})],
            ],
        )
        agent._backends.register("ollama", backend)

        with pytest.raises(MaxIterationsReached):
            await agent.run(session.id, "loop forever")


# ─── run_stream() tests ──────────────────────────────────────────────────────


class TestAgentRunStream:
    @pytest.mark.asyncio
    async def test_run_stream_single_iteration(self, agent, session):
        backend = FakeLLMBackend(responses=["streamed hello"])
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "hi"):
            events.append(type(ev).__name__)

        assert events[-1] == "StreamEnd"
        saved = await agent._sessions.get(session.id)
        assert saved.messages[-1].content == "streamed hello"

    @pytest.mark.asyncio
    async def test_first_message_does_not_force_planning(self, agent, session):
        """Regression: the pre-turn planning gate is gone. Even the FIRST message
        of a session must not run the planner — planning only happens when the
        agent calls the `plan` tool."""
        from navi.core.events import PlanReady, PlanningStatus

        backend = FakeLLMBackend(responses=["streamed hello"])
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "please build the new feature"):
            events.append(ev)

        # No planning events, no plan messages.
        assert not any(isinstance(ev, (PlanningStatus, PlanReady)) for ev in events)
        saved = await agent._sessions.get(session.id)
        assert not any(getattr(m, "is_plan", False) for m in saved.messages)
        # The planner's analysis call never ran (complete() untouched) — only
        # the main loop's stream call hit the backend.
        assert backend._call_idx == 0
        assert backend._stream_idx == 1

    @pytest.mark.asyncio
    async def test_run_stream_emits_model_info(self, agent, session):
        """The agent emits a ModelInfo event carrying the resolved model."""
        from typing import AsyncGenerator

        from navi.llm.base import LLMBackend

        class ModelStampingBackend(LLMBackend):
            async def complete(self, messages, tools=None, temperature=0.7, model=None,
                               think=None, max_tokens=None, **kw):
                raise NotImplementedError

            async def stream_complete(self, messages, tools=None, temperature=0.7,
                                      model=None, think=None, **kw) -> AsyncGenerator[LLMChunk, None]:
                yield LLMChunk(delta="hello", model="resolved-model")
                yield LLMChunk(finish_reason="stop", prompt_tokens=5, completion_tokens=1)

            async def embed(self, texts, model=None):
                return [[0.1] * 768 for _ in texts]

        agent._backends.register("ollama", ModelStampingBackend())
        events = []
        async for ev in agent.run_stream(session.id, "hi"):
            events.append(ev)

        infos = [ev for ev in events if isinstance(ev, ModelInfo)]
        assert len(infos) == 1
        assert infos[0].model == "resolved-model"

    @pytest.mark.asyncio
    async def test_run_stream_emits_model_info_once_per_turn(self, agent, session):
        """ModelInfo is not re-emitted across iterations if the model stays the same."""
        from typing import AsyncGenerator

        from navi.llm.base import LLMBackend

        class ModelStampingBackend(LLMBackend):
            def __init__(self):
                self._call = 0

            async def complete(self, messages, tools=None, temperature=0.7, model=None,
                               think=None, max_tokens=None, **kw):
                raise NotImplementedError

            async def stream_complete(self, messages, tools=None, temperature=0.7,
                                      model=None, think=None, **kw) -> AsyncGenerator[LLMChunk, None]:
                self._call += 1
                if self._call == 1:
                    yield LLMChunk(
                        delta="", model="same-model",
                        finish_reason="tool_calls",
                        tool_calls=[ToolCallRequest(id="1", name="test_tool", arguments={})],
                    )
                else:
                    yield LLMChunk(delta="final answer", model="same-model")
                    yield LLMChunk(finish_reason="stop", prompt_tokens=5, completion_tokens=1)

            async def embed(self, texts, model=None):
                return [[0.1] * 768 for _ in texts]

        agent._backends.register("ollama", ModelStampingBackend())
        events = []
        async for ev in agent.run_stream(session.id, "do something"):
            events.append(ev)

        # Two iterations happened (tool call then final), but ModelInfo fires once.
        infos = [ev for ev in events if isinstance(ev, ModelInfo)]
        assert len(infos) == 1
        assert infos[0].model == "same-model"

    @pytest.mark.asyncio
    async def test_run_stream_tool_calls(self, agent, session):
        backend = FakeLLMBackend(
            responses=["", "final"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
        )
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "do something"):
            events.append(type(ev).__name__)

        assert "ToolStarted" in events
        assert "ToolEvent" in events
        assert events[-1] == "StreamEnd"

    @pytest.mark.asyncio
    async def test_run_stream_emits_todo_updated_after_tool_turn(self, agent, session):
        """A TodoUpdated event is emitted after each tool-execution turn so the
        UI side panel can reflect todo changes from the todo tool."""
        from navi.core.events import TodoUpdated

        backend = FakeLLMBackend(
            responses=["", "final"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
        )
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "do something"):
            events.append(ev)

        todo_events = [ev for ev in events if isinstance(ev, TodoUpdated)]
        assert len(todo_events) >= 1
        # The event carries the session id and a (possibly empty) tasks list.
        assert todo_events[0].session_id == session.id
        assert isinstance(todo_events[0].tasks, list)

    @pytest.mark.asyncio
    async def test_run_stream_stop_event(self, agent, session):
        """Cooperative stop mid-stream yields StreamStopped."""
        from navi.tools._internal.base import current_stop_event

        stop = asyncio.Event()
        token = current_stop_event.set(stop)
        try:
            async def _slow_stream(self, **kwargs):
                yield LLMChunk(delta="a")
                await asyncio.sleep(10)
                yield LLMChunk(delta="b")

            backend = FakeLLMBackend()
            # Monkey-patch stream_complete to be slow
            backend.stream_complete = _slow_stream
            agent._backends.register("ollama", backend)

            stop.set()
            events = []
            async for ev in agent.run_stream(session.id, "hi"):
                events.append(type(ev).__name__)

            assert "StreamStopped" in events
        finally:
            current_stop_event.reset(token)

    @pytest.mark.asyncio
    async def test_run_stream_token_count(self, agent, session):
        backend = FakeLLMBackend(
            responses=["final"],
            prompt_tokens=100,
            completion_tokens=50,
        )
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "hi"):
            if isinstance(ev, StreamEnd):
                events.append(ev)

        assert events[0].token_count == 50
        saved = await agent._sessions.get(session.id)
        assert saved.messages[-1].token_count == 50


# ─── switch_profile reload tests ─────────────────────────────────────────────


class SwitchProfileTool(FakeTool):
    """The real switch_profile's shape: repoint the store, touch nothing else."""

    def __init__(self, store, target):
        super().__init__("switch_profile")
        self._store = store
        self._target = target

    async def execute(self, arguments, ctx=None):
        from navi.tools._internal.base import ToolResult, current_session_id

        sid = (ctx.session_id if ctx else None) or current_session_id.get()
        await self._store.set_profile(sid, self._target)
        return ToolResult(success=True, output=f"Switched to '{self._target}'.")


class ToolRecordingBackend(FakeLLMBackend):
    """Records which tools each LLM call was offered."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.offered: list[set[str]] = []

    async def stream_complete(self, messages, tools=None, **kwargs):
        names = set()
        for tool in tools or []:
            names.add(tool["function"]["name"] if isinstance(tool, dict) else tool.name)
        self.offered.append(names)
        async for chunk in super().stream_complete(messages, tools, **kwargs):
            yield chunk


class TestProfileSwitchReload:
    """A mid-turn switch_profile must rebind the run's tools for the next LLM call.

    Both store paths leave session.profile_id equal to the stored row right after
    a switch: set_profile() rewrites the row (and the run's own save() then put
    its stale value back over it), or it mutates the very Session object the run
    holds, as InMemorySessionStore does. The old guard compared those two after
    the tool batch, so it was never true — the turn kept offering the old
    profile's tools while switch_profile's own answer advertised the new ones.
    """

    def make_agent(self):
        store = InMemorySessionStore()
        profiles = ProfileRegistry()
        profiles.register(
            make_profile("test", enabled_tools=["test_tool", "switch_profile"])
        )
        profiles.register(make_profile("other", enabled_tools=["another_tool"]))
        tools = ToolRegistry()
        tools.register(FakeTool("test_tool"), builtin=True)
        tools.register(FakeTool("another_tool"), builtin=True)
        tools.register(SwitchProfileTool(store, "other"), builtin=True)
        backends = BackendRegistry()
        return Agent(
            session_store=store,
            profile_registry=profiles,
            tool_registry=tools,
            backend_registry=backends,
        )

    @pytest.mark.asyncio
    async def test_the_next_llm_call_is_offered_the_new_profiles_tools(self):
        agent = self.make_agent()
        backend = ToolRecordingBackend(
            responses=["", "final"],
            tool_calls=[
                [
                    ToolCallRequest(
                        id="1",
                        name="switch_profile",
                        arguments={"profile_id": "other"},
                    )
                ],
                None,
            ],
        )
        agent._backends.register("ollama", backend)
        session = await agent._sessions.create(profile_id="test")

        async for _ in agent.run_stream(session.id, "switch please"):
            pass

        assert "test_tool" in backend.offered[0]
        assert backend.offered[1] == {"another_tool"}
        # The switch is what the session ends on, not just what the tool said.
        assert session.profile_id == "other"


# ─── final-turn intercept tests ──────────────────────────────────────────────


class TestFinalIntercept:
    """The final-turn intercept: a bare-text turn (no tool call) with open todo
    steps is continued with a recovery nudge instead of closing the run, up to
    final_intercept_limit. This is the structural patch for the "stopped to
    narrate" failure mode that bypasses the tool-loop nudges (anti-stall,
    adaptive re-plan) because those fire in the *next* iteration's pre_turn,
    which never comes when the turn ends on bare text.
    """

    @pytest_asyncio.fixture
    async def kv(self):
        """In-memory KV so set_tasks persists a plan the agent's has_open_steps
        can read. Mirrors the _fake_kv fixture in tests/unit/tools/test_todo.py."""
        from navi.store import KvStore
        from tests.conftest_factory import FakePool

        class _MemKv(KvStore):
            def __init__(self):
                self._data: dict[tuple, str] = {}

            async def _get_pool(self):
                return FakePool()

            async def get(self, user_id, session_id, scope, key):
                return self._data.get((user_id or "", session_id, scope, key))

            async def set(self, user_id, session_id, scope, key, value):
                self._data[(user_id or "", session_id, scope, key)] = value

            async def get_all(self, user_id, session_id, scope):
                return {
                    k[3]: v
                    for k, v in self._data.items()
                    if k[:3] == (user_id or "", session_id, scope)
                }

            async def delete(self, user_id, session_id, scope, key):
                self._data.pop((user_id or "", session_id, scope, key), None)

            async def clear_scope(self, user_id, session_id, scope):
                keys = [k for k in self._data if k[:3] == (user_id or "", session_id, scope)]
                for k in keys:
                    del self._data[k]

        from navi.tools import todo as _mod

        store = _MemKv()
        _mod._kv_store = store
        yield store
        _mod._kv_store = None

    def _configure(self, agent, *, limit):
        profile = agent._profiles.get("test")
        profile.max_iterations = 20
        profile.final_intercept_enabled = True
        profile.final_intercept_limit = limit

    @pytest.mark.asyncio
    async def test_bare_text_with_open_todo_intercepts_and_continues(self, agent, session, kv):
        """A bare-text turn while the todo has open steps is not closed: the loop
        continues (n-1 intercepts) until the limit is hit, then one final StreamEnd."""
        from navi.tools.todo import set_tasks

        self._configure(agent, limit=2)
        await set_tasks(session.id, ["step a", "step b"])  # open steps -> has_open_steps True
        backend = FakeLLMBackend(
            responses=["I'll try another approach", "I'll try again", "done"],
            tool_calls=[None, None, None],
        )
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "do the task"):
            events.append(ev)

        saved = await agent._sessions.get(session.id)
        assistant = [m for m in saved.messages if m.role == "assistant"]
        # Two intercepts (continue) + one genuine final = three bare-text turns.
        assert len(assistant) == 3
        # Exactly one StreamEnd, at the very end — the loop never closed early.
        stream_ends = [ev for ev in events if isinstance(ev, StreamEnd)]
        assert len(stream_ends) == 1

    @pytest.mark.asyncio
    async def test_intercept_disabled_when_flag_off(self, agent, session, kv):
        """final_intercept_enabled=False: a bare-text turn with open steps
        finalises immediately, no continuation."""
        from navi.tools.todo import set_tasks

        profile = agent._profiles.get("test")
        profile.max_iterations = 20
        profile.final_intercept_enabled = False
        profile.final_intercept_limit = 2
        await set_tasks(session.id, ["step a", "step b"])
        backend = FakeLLMBackend(responses=["I'll try another approach"], tool_calls=[None])
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "do the task"):
            events.append(ev)

        saved = await agent._sessions.get(session.id)
        assistant = [m for m in saved.messages if m.role == "assistant"]
        assert len(assistant) == 1
        assert len([ev for ev in events if isinstance(ev, StreamEnd)]) == 1

    @pytest.mark.asyncio
    async def test_empty_todo_no_intercept(self, agent, session, kv):
        """No plan (empty todo) -> has_open_steps False -> no intercept; the turn
        closes normally. Protects casual messages and observe-only runs."""
        self._configure(agent, limit=2)
        # No set_tasks: the plan is empty.
        backend = FakeLLMBackend(responses=["all done here"], tool_calls=[None])
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "hi"):
            events.append(ev)

        saved = await agent._sessions.get(session.id)
        assistant = [m for m in saved.messages if m.role == "assistant"]
        assert len(assistant) == 1
        assert len([ev for ev in events if isinstance(ev, StreamEnd)]) == 1

    @pytest.mark.asyncio
    async def test_limit_one_intercepts_once(self, agent, session, kv):
        """final_intercept_limit=1: one intercept, then the second bare-text turn
        finalises (limit reached)."""
        from navi.tools.todo import set_tasks

        self._configure(agent, limit=1)
        await set_tasks(session.id, ["step a", "step b"])
        backend = FakeLLMBackend(
            responses=["I'll try again", "done"],
            tool_calls=[None, None],
        )
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "do the task"):
            events.append(ev)

        saved = await agent._sessions.get(session.id)
        assistant = [m for m in saved.messages if m.role == "assistant"]
        assert len(assistant) == 2  # one intercept + one final
        assert len([ev for ev in events if isinstance(ev, StreamEnd)]) == 1

    @pytest.mark.asyncio
    async def test_nudge_text_escalates_and_is_injected_next_iteration(self, agent, session, kv):
        """The 1st-intercept nudge is soft, the 2nd is hard ("second stop"); each is
        injected as a system message into the NEXT iteration's context (not the
        closed turn's). Verified by capturing the messages each LLM call sees."""
        from navi.tools.todo import set_tasks

        self._configure(agent, limit=2)
        await set_tasks(session.id, ["step a", "step b"])

        captured: list[list[Message]] = []

        class CapturingBackend(FakeLLMBackend):
            async def stream_complete(self, messages, **kwargs):
                captured.append(list(messages))
                async for chunk in super().stream_complete(messages, **kwargs):
                    yield chunk

        backend = CapturingBackend(
            responses=["I'll try another approach", "I'll try again", "done"],
            tool_calls=[None, None, None],
        )
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "do the task"):
            events.append(ev)

        # Three LLM calls (2 intercepts + 1 final). captured[0] has no nudge yet;
        # captured[1] carries the soft (1st) nudge; captured[2] the hard (2nd).
        assert len(captured) == 3

        def _sys_content(idx, needle):
            return any(
                m.role == "system" and needle in (m.content or "")
                for m in captured[idx]
            )

        assert not _sys_content(0, "Final-turn check")  # first turn: no nudge yet
        assert _sys_content(1, "Final-turn check")
        assert "second stop" not in "".join(
            (m.content or "") for m in captured[1] if m.role == "system"
        )
        assert _sys_content(2, "second stop")
        # The run still closed exactly once.
        assert len([ev for ev in events if isinstance(ev, StreamEnd)]) == 1


# ─── compact_stream() tests ──────────────────────────────────────────────────


class TestAgentCompactStream:
    """Forced /compact: bypasses the token threshold and runs the real compressor."""

    @pytest_asyncio.fixture
    async def session_with_history(self, agent, session):
        """A session with enough turns that there is something to summarize."""
        # context_keep_recent defaults to 8 turns; add 10 so the oldest turns
        # fall into to_summarize (>= 2 messages) and compress_context does work.
        for i in range(10):
            session.context.append(Message(role="user", content=f"user {i}"))
            session.context.append(Message(role="assistant", content=f"assistant {i}"))
            session.messages.append(Message(role="user", content=f"user {i}"))
            session.messages.append(Message(role="assistant", content=f"assistant {i}"))
        await agent._sessions.save(session)
        return session

    @pytest.mark.asyncio
    async def test_compact_emits_started_then_compressed(self, agent, session_with_history):
        # The agent fixture's FakeLLMBackend returns "hello" — a valid summary.
        events = [ev async for ev in agent.compact_stream(session_with_history.id)]

        assert isinstance(events[0], CompressionStarted)
        assert isinstance(events[-1], ContextCompressed)
        # Compression actually shrank the LLM context (10 turns -> summary + kept recent).
        assert events[-1].messages_after < events[-1].messages_before

    @pytest_asyncio.fixture
    async def session_one_long_turn(self, agent, session):
        """The navi_code shape: a single long autonomous turn — one user message
        followed by many tool/assistant iterations. This is ONE turn, so the
        turn-based partition (len(turns) <= keep_recent) finds nothing to
        summarize. Forced compact must still compress via the intra-turn split
        (keep_recent_messages), otherwise the user always sees
        "Nothing to compact yet — the context is still small" regardless of how
        long the turn grew. See keep_recent_messages in compact_stream."""
        session.context.append(Message(role="user", content="do the task"))
        session.messages.append(Message(role="user", content="do the task"))
        for i in range(20):
            session.context.append(Message(role="assistant", content=f"step {i}"))
            session.messages.append(Message(role="assistant", content=f"step {i}"))
        await agent._sessions.save(session)
        return session

    @pytest.mark.asyncio
    async def test_compact_compresses_single_long_turn(self, agent, session_one_long_turn):
        """Regression: a single long turn (navi_code shape) must compress, not
        report 'nothing to compact'. The intra-turn fallback in compact_stream
        (keep_recent_messages=max(12, context_keep_recent*2)) is what makes this
        work — without it forced compact always returns NothingToCompactError
        once the conversation is a single turn."""
        events = [ev async for ev in agent.compact_stream(session_one_long_turn.id)]

        assert isinstance(events[0], CompressionStarted)
        assert isinstance(events[-1], ContextCompressed)
        assert events[-1].messages_after < events[-1].messages_before

    @pytest.mark.asyncio
    async def test_compact_bypasses_threshold_even_when_context_small(self, agent, session):
        """Forced compact is NOT gated by should_compress — but still needs >= 2
        summarizable messages. A near-empty session raises NothingToCompactError
        (not silently does nothing), which the orchestrator surfaces as feedback."""
        session.context.append(Message(role="user", content="hi"))
        session.context.append(Message(role="assistant", content="hello"))
        await agent._sessions.save(session)

        with pytest.raises(NothingToCompactError):
            async for _ in agent.compact_stream(session.id):
                pass

    @pytest.mark.asyncio
    async def test_compact_raises_when_compression_disabled(self, agent, session_with_history, monkeypatch):
        import navi.core.agent as agent_mod
        import navi.config as config

        # Settings is frozen — swap the whole object the agent module sees.
        disabled = config.Settings(
            database_url=config.settings.database_url,
            context_compression_enabled=False,
        )
        monkeypatch.setattr(agent_mod, "settings", disabled)
        with pytest.raises(NothingToCompactError):
            async for _ in agent.compact_stream(session_with_history.id):
                pass

    @pytest.mark.asyncio
    async def test_compact_unknown_session_raises(self, agent):
        with pytest.raises(SessionNotFound):
            async for _ in agent.compact_stream("no-such-session"):
                pass

    @pytest.mark.asyncio
    async def test_midturn_no_compression_started_when_nothing_to_compress(self, agent, session, monkeypatch):
        """C: even when the token gate fires, CompressionStarted is NOT emitted
        if would_compress says the partition cannot shrink the context — no
        misleading "compression" status with no work done. The gate is forced
        True here so we exercise the would_compress guard directly."""
        import navi.core.agent as agent_mod
        from navi.core.events import CompressionStarted

        # Small context that partition cannot shrink (would_compress False).
        session.context.append(Message(role="user", content="hi"))
        session.context.append(Message(role="assistant", content="hello"))
        await agent._sessions.save(session)

        # Force the token gate open; without this the small context would not
        # reach the would_compress check at all.
        monkeypatch.setattr(agent_mod, "should_compress", lambda *a, **k: True)

        events = [
            ev
            async for ev in agent._compression_events_midturn(
                session,
                llm=agent._get_backend("ollama"),
                profile=agent._profiles.get("test"),
                session_id=session.id,
                iteration=1,
                ctx_injections=[],
                mem=None,
            )
        ]
        assert not any(isinstance(ev, CompressionStarted) for ev in events)


# ─── run_ephemeral() tests ───────────────────────────────────────────────────


class TestAgentRunEphemeral:
    @pytest.mark.asyncio
    async def test_run_ephemeral_complete(self, agent):
        backend = FakeLLMBackend(responses=["subagent result"])
        agent._backends.register("ollama", backend)

        result, ok = await agent.run_ephemeral("task", profile_id="test")
        assert "subagent result" in result
        assert "[Sub-agent stopped: completed]" in result
        assert ok is True

    @pytest.mark.asyncio
    async def test_run_ephemeral_max_iterations(self, agent):
        backend = FakeLLMBackend(
            responses=[""],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
            ],
        )
        agent._backends.register("ollama", backend)

        result, ok = await agent.run_ephemeral(
            "task", profile_id="test", max_iterations=1
        )
        assert ok is False
        assert "iteration limit" in result.lower()

    @pytest.mark.skip(reason="run_ephemeral uses 'import time as _time' inside the function; CPython LOAD_GLOBAL caching makes module-level mock replacement unreliable in pytest-asyncio.")
    @pytest.mark.asyncio
    async def test_run_ephemeral_timeout(self, agent):
        pass

    @pytest.mark.asyncio
    async def test_run_ephemeral_planning_tokens_accumulated(self, agent):
        """Planning phase AIHelperTokensUsed contributes to SubagentComplete."""
        from navi.core.events import AIHelperTokensUsed
        from navi.tools._internal.base import current_event_sink

        backend = FakeLLMBackend(responses=["final"])
        agent._backends.register("ollama", backend)

        # Force planning by setting subagent_planning_enabled on profile
        profile = agent._profiles.get("test")
        profile.subagent_planning_enabled = True

        sink = asyncio.Queue()
        token = current_event_sink.set(sink)
        try:
            # Mock planning to emit AIHelperTokensUsed
            original_planning_run = agent._planning.run

            async def _mock_planning(*args, **kwargs):
                yield AIHelperTokensUsed(prompt_tokens=5, completion_tokens=10)
                yield AIHelperTokensUsed(prompt_tokens=3, completion_tokens=7)

            agent._planning.run = _mock_planning

            result, ok = await agent.run_ephemeral("task", profile_id="test")
            assert ok is True

            # Drain sink for SubagentComplete
            subagent_complete = None
            while not sink.empty():
                item = await sink.get()
                if isinstance(item, SubagentComplete):
                    subagent_complete = item

            # Planning completion tokens: 10 + 7 = 17
            # Final LLM call: 0 (no tokens in FakeLLMBackend default)
            assert subagent_complete is not None
            assert subagent_complete.token_count == 17
        finally:
            current_event_sink.reset(token)
            agent._planning.run = original_planning_run

    @pytest.mark.asyncio
    async def test_run_ephemeral_thinking_stall(self, agent):
        """Subagent that produces only thinking for too long is aborted."""
        async def _thinking_only(self, **kwargs):
            for _ in range(200):
                yield LLMChunk(thinking="thinking " * 100)
            yield LLMChunk(delta="done", finish_reason="stop")

        backend = FakeLLMBackend()
        backend.stream_complete = _thinking_only
        agent._backends.register("ollama", backend)

        result, ok = await agent.run_ephemeral("task", profile_id="test")
        assert ok is False
        assert "thinking" in result.lower() or "stall" in result.lower()


class _SnapshotSessionStore(InMemorySessionStore):
    """Session store that mimics a real DB boundary.

    ``save()`` deep-copies the session into a snapshot; ``get()`` returns the last
    saved snapshot (a fresh copy), not the live object the agent is mutating. So
    mutations that were never ``save()``d are invisible to ``get()`` — exactly like
    ``PgSessionStore`` where messages with ``sequence_number < 0`` only persist on
    ``save()``. This lets us assert that a crash mid-turn does not lose work that
    was incrementally flushed (B1).
    """

    def __init__(self) -> None:
        super().__init__()
        self._snapshots: dict[str, Session] = {}

    async def save(self, session: Session) -> None:
        session.last_active = datetime.now(timezone.utc)
        # Mirror PgSessionStore: assign sequence numbers to new messages on save.
        for m in session.messages:
            if m.sequence_number < 0:
                m.sequence_number = session.db_next_sequence
                session.db_next_sequence += 1
        self._snapshots[session.id] = copy.deepcopy(session)
        self._sessions[session.id] = session

    async def get(self, session_id: str) -> Session | None:
        snap = self._snapshots.get(session_id)
        if snap is not None:
            return copy.deepcopy(snap)
        return self._sessions.get(session_id)


class _CrashBackend(FakeLLMBackend):
    """FakeLLMBackend whose ``stream_complete`` raises on the Nth call.

    Simulates a server crash / CancelledError arriving mid-turn, so we can assert
    that tool work completed before the crash was already persisted (B1).
    """

    def __init__(
        self,
        responses: list[str],
        tool_calls: list[list[ToolCallRequest] | None],
        raise_on_call: int,
        exc: BaseException = RuntimeError("simulated crash"),
    ) -> None:
        super().__init__(responses=responses, tool_calls=tool_calls)
        self._raise_on_call = raise_on_call
        self._exc = exc
        self._sc_idx = 0

    async def stream_complete(self, messages, tools=None, temperature=0.7, model=None,
                              think=None, **kwargs):
        idx = self._sc_idx
        self._sc_idx += 1
        if idx == self._raise_on_call:
            raise self._exc
        content = self._responses[idx] if idx < len(self._responses) else ""
        tcalls = self._tool_calls[idx] if idx < len(self._tool_calls) else None
        if content:
            yield LLMChunk(delta=content)
        yield LLMChunk(
            finish_reason="tool_calls" if tcalls else "stop",
            tool_calls=tcalls,
        )


def _build_persistence_agent(store: InMemorySessionStore, backend) -> Agent:
    profiles = ProfileRegistry()
    profile = make_profile("test")
    profiles.register(profile)
    tools = make_registry_with_tools()
    backends = BackendRegistry()
    backends.register("ollama", backend)
    return Agent(
        session_store=store,
        profile_registry=profiles,
        tool_registry=tools,
        backend_registry=backends,
    )


class TestAgentPersistenceOnCrash:
    @pytest.mark.asyncio
    async def test_tool_results_persisted_on_crash_midturn(self):
        """A crash on a later LLM call must not lose tool results already completed
        in earlier iterations — B1 persists each tool result as it completes."""
        store = _SnapshotSessionStore()
        backend = _CrashBackend(
            responses=["", "unused"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
            raise_on_call=1,
        )
        agent = _build_persistence_agent(store, backend)
        sess = await agent._sessions.create(profile_id="test")

        with pytest.raises(RuntimeError, match="simulated crash"):
            async for _ in agent.run_stream(sess.id, "do work"):
                pass

        saved = await agent._sessions.get(sess.id)
        roles = [m.role for m in saved.messages]
        # Iteration 0's assistant tool-call decision + tool result were flushed by
        # B1 before the crash on iteration 1. Without B1 only the user message
        # (saved at turn start) would survive.
        assert "tool" in roles, f"tool result missing from saved state: {roles}"
        assert any(
            m.role == "assistant" and m.tool_calls for m in saved.messages
        ), f"assistant tool-call msg missing: {roles}"

    @pytest.mark.asyncio
    async def test_cancel_flushes_completed_tool_work(self):
        """CancelledError mid-turn (server restart/shutdown) must not lose tool
        results already completed — B1 flushed them incrementally before the cancel."""
        store = _SnapshotSessionStore()
        backend = _CrashBackend(
            responses=["", "unused"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
            raise_on_call=1,
            exc=asyncio.CancelledError(),
        )
        agent = _build_persistence_agent(store, backend)
        sess = await agent._sessions.create(profile_id="test")

        with pytest.raises(asyncio.CancelledError):
            async for _ in agent.run_stream(sess.id, "do work"):
                pass

        saved = await agent._sessions.get(sess.id)
        roles = [m.role for m in saved.messages]
        assert "tool" in roles, f"tool result missing after cancel: {roles}"