"""Unit tests for the parallel tool-call batch (Ф3)."""

import asyncio
import time
from types import SimpleNamespace

import pytest

from navi.core.agent import Agent, AgentTurnContext
from navi.core.events import TextDelta, ToolEvent, ToolStarted
from navi.core.tool_executor import ToolExecutor
from navi.llm.base import ToolCallRequest
from navi.tools._internal.base import ToolResult, current_event_sink


class FakeStore:
    def __init__(self):
        self.saves = 0

    async def save(self, session):
        self.saves += 1


class SlowTool:
    """Sleeps, optionally emits a live event, returns output."""

    def __init__(self, name, delay, emit=None, fail=False):
        self.name = name
        self._delay = delay
        self._emit = emit
        self._fail = fail

    async def execute(self, arguments, ctx=None):
        await asyncio.sleep(self._delay)
        if self._emit is not None:
            sink = current_event_sink.get()
            await sink.put(self._emit)
        if self._fail:
            raise ValueError(f"{self.name} blew up")
        return ToolResult(success=True, output=f"done {self.name}")


def make_agent():
    agent = object.__new__(Agent)
    agent._tool_executor = ToolExecutor(SimpleNamespace(_middlewares=[]))
    agent._sessions = FakeStore()
    return agent, agent._sessions


def make_tcs(*names):
    return [ToolCallRequest(id=f"tc-{i}", name=n, arguments={})
            for i, n in enumerate(names)]


def make_turn_ctx(parallel=True):
    return AgentTurnContext(turn_start=time.monotonic(), parallel_tool_calls=parallel)


class TestParallelBatch:
    async def test_wall_clock_shorter_than_sequential(self):
        agent, _ = make_agent()
        tools = [SlowTool("a", 0.25), SlowTool("b", 0.25)]
        tool_map = {t.name: t for t in tools}

        async def run():
            return [ev async for ev in agent._execute_tools_parallel(
                make_tcs("a", "b"), tools, make_turn_ctx(), SimpleNamespace(messages=[], context=[]),
                None, None)]

        start = time.monotonic()
        events = await run()
        elapsed = time.monotonic() - start
        assert elapsed < 0.45  # sequential would take ≥0.5s
        tool_events = [e for e in events if isinstance(e, ToolEvent)]
        assert len(tool_events) == 2

    async def test_tool_started_all_before_any_result(self):
        agent, _ = make_agent()
        tools = [SlowTool("a", 0.25), SlowTool("b", 0.25)]
        events = []
        async for ev in agent._execute_tools_parallel(
            make_tcs("a", "b"), tools, make_turn_ctx(),
            SimpleNamespace(messages=[], context=[]), None, None,
        ):
            events.append(ev)
        started = [i for i, e in enumerate(events) if isinstance(e, ToolStarted)]
        done = [i for i, e in enumerate(events) if isinstance(e, ToolEvent)]
        assert started == [0, 1]
        assert done == [2, 3]

    async def test_results_in_call_order_despite_finish_order(self):
        agent, _ = make_agent()
        # second tool finishes first
        tools = [SlowTool("slow", 0.2), SlowTool("fast", 0.01)]
        events = [ev async for ev in agent._execute_tools_parallel(
            make_tcs("slow", "fast"), tools, make_turn_ctx(),
            SimpleNamespace(messages=[], context=[]), None, None)]
        tool_events = [e for e in events if isinstance(e, ToolEvent)]
        assert [e.tool_name for e in tool_events] == ["slow", "fast"]

    async def test_live_events_merged_through_shared_queue(self):
        agent, _ = make_agent()
        tools = [SlowTool("a", 0.05, emit=TextDelta(delta="A")),
                 SlowTool("b", 0.05, emit=TextDelta(delta="B"))]
        events = [ev async for ev in agent._execute_tools_parallel(
            make_tcs("a", "b"), tools, make_turn_ctx(),
            SimpleNamespace(messages=[], context=[]), None, None)]
        deltas = [e.delta for e in events if isinstance(e, TextDelta)]
        assert sorted(deltas) == ["A", "B"]

    async def test_one_failure_does_not_kill_neighbours(self):
        agent, store = make_agent()
        tools = [SlowTool("bad", 0.01, fail=True), SlowTool("good", 0.01)]
        session = SimpleNamespace(messages=[], context=[])
        events = [ev async for ev in agent._execute_tools_parallel(
            make_tcs("bad", "good"), tools, make_turn_ctx(), session, None, None)]
        tool_events = [e for e in events if isinstance(e, ToolEvent)]
        by_name = {e.tool_name: e for e in tool_events}
        assert by_name["bad"].success is False
        assert by_name["good"].success is True
        # both tool messages recorded, in call order
        assert [m.name for m in session.messages] == ["bad", "good"]

    async def test_stop_mid_batch_synthesises_results_in_order(self):
        agent, store = make_agent()
        tools = [SlowTool("a", 5.0), SlowTool("b", 5.0)]
        session = SimpleNamespace(messages=[], context=[])
        stop_event = asyncio.Event()

        async def set_stop():
            await asyncio.sleep(0.1)
            stop_event.set()

        asyncio.create_task(set_stop())
        events = [ev async for ev in agent._execute_tools_parallel(
            make_tcs("a", "b"), tools, make_turn_ctx(), session, stop_event, None)]
        tool_events = [e for e in events if isinstance(e, ToolEvent)]
        assert len(tool_events) == 2
        assert all(not e.success for e in tool_events)
        assert all("stopped by the user" in e.result for e in tool_events)
        assert all(m.is_context is False for m in session.messages)

    async def test_single_save_per_batch(self):
        agent, store = make_agent()
        tools = [SlowTool("a", 0.01), SlowTool("b", 0.01)]
        session = SimpleNamespace(messages=[], context=[])
        async for _ in agent._execute_tools_parallel(
            make_tcs("a", "b"), tools, make_turn_ctx(), session, None, None):
            pass
        assert store.saves == 1

    async def test_tool_call_count_incremented(self):
        agent, _ = make_agent()
        tools = [SlowTool("a", 0.01), SlowTool("b", 0.01)]
        turn_ctx = make_turn_ctx()
        async for _ in agent._execute_tools_parallel(
            make_tcs("a", "b"), tools, turn_ctx,
            SimpleNamespace(messages=[], context=[]), None, None):
            pass
        assert turn_ctx.tool_call_count == 2


class TestDispatch:
    async def test_gate_off_uses_sequential_path(self):
        """parallel_tool_calls=False keeps the strict sequential ordering."""
        agent, _ = make_agent()
        tools = [SlowTool("a", 0.05), SlowTool("b", 0.05)]
        events = [ev async for ev in agent._execute_tools_with_sink(
            make_tcs("a", "b"), tools, make_turn_ctx(parallel=False),
            SimpleNamespace(messages=[], context=[]), None, None)]
        kinds = [type(e).__name__ for e in events]
        # sequential: Started(a), Event(a), Started(b), Event(b)
        assert kinds == ["ToolStarted", "ToolEvent", "ToolStarted", "ToolEvent"]

    async def test_gate_on_routes_to_parallel(self):
        agent, _ = make_agent()
        tools = [SlowTool("a", 0.25), SlowTool("b", 0.25)]
        events = []
        async for ev in agent._execute_tools_with_sink(
            make_tcs("a", "b"), tools, make_turn_ctx(parallel=True),
            SimpleNamespace(messages=[], context=[]), None, None):
            events.append(ev)
        kinds = [type(e).__name__ for e in events]
        assert kinds[:2] == ["ToolStarted", "ToolStarted"]

    async def test_single_call_batch_stays_sequential(self):
        """One tool call never needs the parallel machinery."""
        agent, _ = make_agent()
        tools = [SlowTool("a", 0.01)]
        events = [ev async for ev in agent._execute_tools_with_sink(
            make_tcs("a"), tools, make_turn_ctx(parallel=True),
            SimpleNamespace(messages=[], context=[]), None, None)]
        kinds = [type(e).__name__ for e in events]
        assert kinds == ["ToolStarted", "ToolEvent"]


class TestRepair:
    def test_repair_dangling_tool_calls(self):
        from navi.core.pg_session_store import repair_dangling_tool_calls
        from navi.llm.base import Message

        assistant = Message(role="assistant", content=None,
                            tool_calls=[ToolCallRequest(id="t1", name="a", arguments={}),
                                        ToolCallRequest(id="t2", name="b", arguments={})])
        answered = Message(role="tool", content="ok", tool_call_id="t1", name="a")
        later = Message(role="user", content="next")
        messages = [assistant, answered, later]

        added = repair_dangling_tool_calls(messages)
        assert added == 1
        # placeholder sits after the answered tool message, before the user msg
        assert messages[2].role == "tool"
        assert messages[2].tool_call_id == "t2"
        assert messages[2].content == "[Interrupted before result]"
        assert messages[2].is_context is False
        assert messages[3] is later

    def test_repair_noop_when_complete(self):
        from navi.core.pg_session_store import repair_dangling_tool_calls
        from navi.llm.base import Message

        messages = [
            Message(role="assistant", content=None,
                    tool_calls=[ToolCallRequest(id="t1", name="a", arguments={})]),
            Message(role="tool", content="ok", tool_call_id="t1", name="a"),
        ]
        assert repair_dangling_tool_calls(messages) == 0
        assert len(messages) == 2