"""Tool execution: scope resolution, workers, parallel and sequential batches."""

import asyncio
from typing import TYPE_CHECKING

import structlog

from navi.config import settings
from navi.exceptions import (
    LLMBackendError,
    LLMConnectionError,
)
from navi.llm.base import LLMBackend, Message
from navi.tools._internal.base import (
    current_event_sink,
    current_user_role,
)

from ..agent_run_context import AgentTurnContext
from ..events import (
    AgentEvent,
    AIHelperTokensUsed,
    ProfileSwitched,
    SubagentComplete,
    TodoUpdated,
    ToolEvent,
    ToolStarted,
)

if TYPE_CHECKING:
    from navi.profiles.base import AgentProfile


log = structlog.get_logger()


from .helpers import _TOOL_DONE, _TaggedSink


class AgentToolsMixin:
    async def _apply_mcp_scope(self, session):
        """Restrict this run's MCP tool list to the servers its user may use.

        A BYOK-gated server with no personal key must not appear in the list at
        all: the manager refuses the call anyway, and a tool the model can see
        but never call is read as a transient fault and retried. Admins (and the
        no-auth local user, whose role is ``admin``) get ``None`` — the shared
        credential in the config is theirs, so nothing is filtered.

        Also records *which* servers were withheld for want of a key. The reason
        is not stored anywhere else and is not derivable by the caller: it is the
        complement of what ``visible_servers`` kept, so it is computed from that
        same value rather than read a second time (a second read of the key store
        could disagree with the first).

        Returns the pair of ContextVar tokens; callers reset them where they
        already reset the session and cwd vars.
        """
        from navi.tools._internal.base import (
            current_allowed_mcp_servers,
            current_keyless_mcp_servers,
            current_user_id,
        )

        if self._mcp_manager is None:
            allowed, keyless = None, frozenset()
        else:
            try:
                allowed = await self._mcp_manager.visible_servers(
                    current_user_role.get(), current_user_id.get() or session.user_id
                )
                # Withheld = gated servers the run did not get to keep. None means
                # "no restriction", so nothing was withheld.
                keyless = (
                    frozenset()
                    if allowed is None
                    else frozenset(self._mcp_manager.byok_servers() - allowed)
                )
            except Exception:
                # Fail closed: an unanswerable visibility question must not fall
                # back to "everything is available". Every gated server then counts
                # as withheld, which matches the empty tool list the model gets.
                log.warning("agent.mcp_scope_failed", exc_info=True)
                allowed = frozenset()
                keyless = frozenset(self._mcp_manager.byok_servers())

        return (
            current_allowed_mcp_servers.set(allowed),
            current_keyless_mcp_servers.set(keyless),
        )

    async def _run_workers(
        self,
        session,
        llm: LLMBackend,
        model: str,
        context_tokens: int | None,
        profile: "AgentProfile | None" = None,
    ) -> list[AgentEvent]:
        """Run all workers sequentially; collect their events."""
        from navi.workers.base import WorkerContext

        ctx = WorkerContext(
            session_id=session.id,
            context_tokens=context_tokens,
            max_context_tokens=self._window_for(profile),
            llm=llm,
            model=model,
            temperature=settings.context_summary_temperature,
            session_store=self._sessions,
            profile=profile,
            # The backend cut the last prompt down (record_real_baseline saw it).
            # Its token count is below the trigger, but the context it came from
            # is over the window — compress now instead of waiting for a
            # threshold that the truncated count will never reach.
            context_truncated=self._compressor.pop_context_truncated(),
            memory_store=self._memory,
        )
        events: list[AgentEvent] = []
        for worker in self._workers:
            try:
                result = await worker.run(session, ctx)
                events.extend(result.events)
            except Exception:
                log.warning("agent.worker_failed", worker=type(worker).__name__, exc_info=True)
        return events

    async def _execute_tools_with_sink(
        self, turn_tool_calls, tools, turn_ctx: AgentTurnContext, session, stop_event, tool_ctx=None
    ):
        """Execute tool calls, sequentially or as a parallel batch.

        The mode is decided per-turn by ``turn_ctx.parallel_tool_calls``
        (profile override ? global setting, resolved at tool_ctx build).
        """
        # A profile switch replaces the toolset mid-turn, so it must not share a
        # batch: every neighbour would still be dispatched against the old
        # tool_map and die with "tool 'X' not found" (seen with reload_tools called
        # just after a switch). Run it alone first, then re-resolve the tools
        # it switched to for the rest of the batch. The turn's own ``profile``
        # binding stays old on purpose — the end-of-iteration reload in
        # run_stream() is what rebinds profile/llm/schemas.
        switching = next((tc for tc in turn_tool_calls if tc.name == "switch_profile"), None)
        if switching is not None and len(turn_tool_calls) > 1:
            rest = [tc for tc in turn_tool_calls if tc is not switching]
            switched_to = None
            async for ev in self._execute_tools_sequential(
                [switching], tools, turn_ctx, session, stop_event, tool_ctx
            ):
                if isinstance(ev, ProfileSwitched):
                    switched_to = ev.profile_id
                yield ev
            if switched_to:
                new_tools = self._tool_list_for_profile(switched_to)
                if new_tools is not None:
                    tools = new_tools
            turn_tool_calls = rest

        if turn_tool_calls and turn_ctx.parallel_tool_calls and len(turn_tool_calls) > 1:
            async for ev in self._execute_tools_parallel(
                turn_tool_calls, tools, turn_ctx, session, stop_event, tool_ctx
            ):
                yield ev
        else:
            async for ev in self._execute_tools_sequential(
                turn_tool_calls, tools, turn_ctx, session, stop_event, tool_ctx
            ):
                yield ev

    async def _execute_tools_parallel(
        self, turn_tool_calls, tools, turn_ctx: AgentTurnContext, session, stop_event, tool_ctx=None
    ):
        """Execute an independent batch of tool calls concurrently.

        ToolStarted is emitted for every call up-front (in call order) so the
        UI shows the whole batch immediately. Live events from all tools merge
        through one tagged queue; results and tool messages are appended in
        call order after the batch settles, with a single session save.

        Only batch independent calls — two actions over the same terminal in
        one batch race each other.
        """
        from navi.tools.todo import started_metadata_for_call

        tool_map = {t.name: t for t in tools}
        for tc in turn_tool_calls:
            yield ToolStarted(
                tool_name=tc.name,
                arguments=tc.arguments,
                tool_call_id=tc.id,
                metadata=await started_metadata_for_call(tc, tool_ctx),
            )

        shared: asyncio.Queue = asyncio.Queue()
        holders: list[list] = [[] for _ in turn_tool_calls]
        tasks: list[asyncio.Task] = []

        for idx, tc in enumerate(turn_tool_calls):
            tagged = _TaggedSink(shared, idx)
            sink_token = current_event_sink.set(tagged)

            async def _run_with_sentinel(_tc=tc, _idx=idx, _holder=holders[idx], _sink=tagged):
                try:
                    _holder.append(
                        await self._tool_executor._run_single_tool(_tc, tool_map, ctx=tool_ctx)
                    )
                except Exception as exc:
                    _holder.append(exc)
                finally:
                    await _sink.put(_TOOL_DONE)

            tasks.append(asyncio.create_task(_run_with_sentinel()))
            current_event_sink.reset(sink_token)

        try:
            pending = set(range(len(tasks)))
            stopped = False
            while pending:
                try:
                    item = await asyncio.wait_for(shared.get(), timeout=1.0)
                except asyncio.TimeoutError:
                    if stop_event and stop_event.is_set():
                        stopped = True
                        break
                    continue

                _idx, value = item
                if value is _TOOL_DONE:
                    pending.discard(_idx)
                    continue
                if isinstance(value, SubagentComplete):
                    turn_ctx.subagent_tokens += value.token_count
                    turn_ctx.tool_call_count += value.tool_call_count
                elif isinstance(value, AIHelperTokensUsed):
                    turn_ctx.subagent_tokens += value.completion_tokens
                else:
                    yield value

            if stopped:
                # Stop must actually stop work: cancel in-flight tools the same
                # way the sequential path does (docs/tasks.md: stop mid-batch
                # cancels every unfinished tool). Calls that already finished
                # keep their real results — only unfinished ones get the
                # synthetic stopped note below; erasing a completed
                # backgrounded call (task_id) would orphan the task in the
                # client UI.
                log.info(
                    "agent.batch_stopped",
                    tools=[tc.name for tc in turn_tool_calls],
                )
                for t in tasks:
                    if not t.done():
                        t.cancel()
                await asyncio.gather(*tasks, return_exceptions=True)
                # Same reason as in the sequential path: whatever the cancelled
                # tools announced before unwinding still belongs to cards the
                # client is showing, so it is forwarded rather than dropped.
                while True:
                    try:
                        _idx, value = shared.get_nowait()
                    except asyncio.QueueEmpty:
                        break
                    if value is _TOOL_DONE:
                        continue
                    if isinstance(value, SubagentComplete):
                        turn_ctx.subagent_tokens += value.token_count
                        turn_ctx.tool_call_count += value.tool_call_count
                    elif isinstance(value, AIHelperTokensUsed):
                        turn_ctx.subagent_tokens += value.completion_tokens
                    else:
                        yield value

            # Results and messages in call order, one save for the batch.
            # A stopped batch mixes real results (finished calls) with the
            # synthetic stopped note (cancelled calls).
            for tc, holder in zip(turn_tool_calls, holders):
                if not holder:
                    if stopped:
                        yield ToolEvent(
                            tool_name=tc.name,
                            arguments=tc.arguments,
                            result="Tool execution was stopped by the user.",
                            success=False,
                            tool_call_id=tc.id,
                        )
                        session.messages.append(
                            Message(
                                role="tool",
                                content="Tool execution was stopped by the user.",
                                tool_call_id=tc.id,
                                name=tc.name,
                                metadata={},
                                is_context=False,
                            )
                        )
                        continue
                    r = RuntimeError("tool task produced no result")
                else:
                    r = holder[0]
                if isinstance(r, Exception):
                    if isinstance(r, (LLMBackendError, LLMConnectionError)):
                        raise r
                    log.warning("agent.tool_exception", tool=tc.name, error=str(r))
                    tool_event = ToolEvent(
                        tool_name=tc.name,
                        arguments=tc.arguments,
                        result=f"Error: {r}",
                        success=False,
                        tool_call_id=tc.id,
                    )
                    msg = Message(
                        role="tool",
                        content=f"Error: {r}",
                        tool_call_id=tc.id,
                        name=tc.name,
                        metadata={},
                    )
                    image_msg = None
                else:
                    tool_event, msg, image_msg = r

                turn_ctx.tool_call_count += 1
                yield tool_event
                session.messages.append(msg)
                session.context.append(msg)
                if image_msg:
                    session.messages.append(image_msg)
                    session.context.append(image_msg)
            await self._sessions.save(session)
        finally:
            # Cancel whatever is still running (generator closed mid-batch or
            # an outer cancellation of the run), then gather to contain each
            # cancelled task — awaiting the results directly would re-raise
            # CancelledError past the results already produced, while an
            # outer cancellation of the run itself still propagates.
            for t in tasks:
                if not t.done():
                    t.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)

    async def _execute_tools_sequential(
        self, turn_tool_calls, tools, turn_ctx: AgentTurnContext, session, stop_event, tool_ctx=None
    ):
        """Execute tool calls one at a time with cooperative stop support.

        Polls *stop_event* every second while draining the event sink so the
        Stop button works even during long-running tools (terminal, SSH,
        web search, sub-agent spawn).
        """
        tool_map = {t.name: t for t in tools}
        for tc in turn_tool_calls:
            from navi.tools.todo import started_metadata_for_call

            yield ToolStarted(
                tool_name=tc.name,
                arguments=tc.arguments,
                tool_call_id=tc.id,
                metadata=await started_metadata_for_call(tc, tool_ctx),
            )

            sink: asyncio.Queue = asyncio.Queue()
            sink_token = current_event_sink.set(sink)
            result_holder: list = []

            async def _run_with_sentinel(_tc=tc, _holder=result_holder, _sink=sink):
                try:
                    _holder.append(
                        await self._tool_executor._run_single_tool(_tc, tool_map, ctx=tool_ctx)
                    )
                except Exception as exc:
                    _holder.append(exc)
                finally:
                    await _sink.put(_TOOL_DONE)

            tool_task = asyncio.create_task(_run_with_sentinel())
            current_event_sink.reset(sink_token)

            try:
                # Poll sink with 1s timeout so stop_event is checked during long tools
                stopped = False
                while True:
                    try:
                        item = await asyncio.wait_for(sink.get(), timeout=1.0)
                    except asyncio.TimeoutError:
                        if stop_event and stop_event.is_set():
                            stopped = True
                            if not tool_task.done():
                                tool_task.cancel()
                            break
                        continue

                    if item is _TOOL_DONE:
                        break
                    if isinstance(item, SubagentComplete):
                        turn_ctx.subagent_tokens += item.token_count
                        turn_ctx.tool_call_count += item.tool_call_count
                    elif isinstance(item, AIHelperTokensUsed):
                        turn_ctx.subagent_tokens += item.completion_tokens
                    else:
                        yield item

                if stopped:
                    # A cancelled tool may still have announced work the client
                    # is already showing: a sub-agent emits a closing ToolEvent
                    # for its in-flight tool from the CancelledError handler.
                    # Read those before dropping the queue, or the client keeps
                    # a spinner it can never resolve until the page reloads.
                    await asyncio.gather(tool_task, return_exceptions=True)
                    while True:
                        try:
                            item = sink.get_nowait()
                        except asyncio.QueueEmpty:
                            break
                        if item is _TOOL_DONE:
                            continue
                        if isinstance(item, SubagentComplete):
                            turn_ctx.subagent_tokens += item.token_count
                            turn_ctx.tool_call_count += item.tool_call_count
                        elif isinstance(item, AIHelperTokensUsed):
                            turn_ctx.subagent_tokens += item.completion_tokens
                        else:
                            yield item
                    # Tool was interrupted by user — record a synthetic cancellation result
                    log.info("agent.tool_stopped", tool=tc.name)
                    yield ToolEvent(
                        tool_name=tc.name,
                        arguments=tc.arguments,
                        result="Tool execution was stopped by the user.",
                        success=False,
                        tool_call_id=tc.id,
                    )
                    session.messages.append(
                        Message(
                            role="tool",
                            content="Tool execution was stopped by the user.",
                            tool_call_id=tc.id,
                            name=tc.name,
                            metadata={},
                            is_context=False,
                        )
                    )
                    await self._sessions.save(session)
                    return

                r = (
                    result_holder[0]
                    if result_holder
                    else RuntimeError("tool task produced no result")
                )
                if isinstance(r, Exception):
                    if isinstance(r, (LLMBackendError, LLMConnectionError)):
                        raise r
                    log.warning("agent.tool_exception", tool=tc.name, error=str(r))
                    tool_event = ToolEvent(
                        tool_name=tc.name,
                        arguments=tc.arguments,
                        result=f"Error: {r}",
                        success=False,
                        tool_call_id=tc.id,
                    )
                    msg = Message(
                        role="tool",
                        content=f"Error: {r}",
                        tool_call_id=tc.id,
                        name=tc.name,
                        metadata={},
                    )
                    image_msg = None
                else:
                    tool_event, msg, image_msg = r

                turn_ctx.tool_call_count += 1
                yield tool_event
                session.messages.append(msg)
                session.context.append(msg)
                if image_msg:
                    session.messages.append(image_msg)
                    session.context.append(image_msg)
                # Persist each tool result as it completes, so a crash /
                # CancelledError (server restart) between tool calls cannot lose
                # already-finished work — only the in-flight single tool is at risk.
                await self._sessions.save(session)
            finally:
                if not tool_task.done():
                    tool_task.cancel()
                # gather(return_exceptions=True) contains the cancelled tool
                # task (awaiting it directly would re-raise CancelledError and
                # abort the turn before the synthetic stopped results land),
                # while an outer cancellation of the run itself propagates.
                await asyncio.gather(tool_task, return_exceptions=True)

    async def _todo_updated(self, session) -> TodoUpdated:
        """Build a TodoUpdated event from the session's current todo state.

        Reads via explicit (user_id, session_id) so it works regardless of the
        ContextVar state at the call site. Sub-agent todos are isolated in their
        own KV row and never reach here (the parent's session.id is the parent's).
        """
        from navi.tools.todo import load_tasks_for

        tasks = await load_tasks_for(session.id, session.user_id)
        return TodoUpdated(session_id=session.id, tasks=tasks)