"""Tool execution helpers — extracted from agent.py."""

import asyncio
import json
from typing import TYPE_CHECKING

import structlog

from navi.llm.base import Message, ToolCallRequest
from navi.tools._internal.base import Tool, ToolResult, current_session_id

if TYPE_CHECKING:
    from navi.core.events import ToolEvent

log = structlog.get_logger()

# Tools whose foreground default timeout is too short for a detached run.
_TIMEOUT_TOOLS = {"terminal", "code_exec", "ssh_exec"}
_BG_TIMEOUT = 300


def _backgroundable_tools() -> set[str]:
    from navi.config import settings

    return {t.strip() for t in settings.backgroundable_tools.split(",") if t.strip()}


def _resolve_tool(tool_map: dict[str, Tool], name: str) -> tuple[str, Tool | None]:
    """Resolve exact tool names plus common MCP alias mistakes."""
    from navi.mcp.tools import is_mcp_tool, parse_mcp_name

    tool = tool_map.get(name)
    if tool is not None:
        return name, tool

    # Support bare tool name when the full MCP name ends with it
    # e.g. "web_search" -> "mcp__navi-web__web_search"
    bare_matches = [
        (candidate_name, candidate)
        for candidate_name, candidate in tool_map.items()
        if is_mcp_tool(candidate_name) and candidate_name.endswith(f"__{name}")
    ]
    if len(bare_matches) == 1:
        return bare_matches[0]

    # Normalized variant (dash vs underscore)
    normalized = name.replace("-", "_")
    normalized_matches = [
        (candidate_name, candidate)
        for candidate_name, candidate in tool_map.items()
        if is_mcp_tool(candidate_name) and candidate_name.replace("-", "_") == normalized
    ]
    if len(normalized_matches) == 1:
        return normalized_matches[0]

    # Fallback: old underscore format like mcp_server_tool -> mcp__server__tool
    old_format_matches = [
        (candidate_name, candidate)
        for candidate_name, candidate in tool_map.items()
        if is_mcp_tool(candidate_name) and name.startswith("mcp_")
    ]
    for candidate_name, candidate in old_format_matches:
        parsed = parse_mcp_name(candidate_name)
        if parsed is None:
            continue
        server_name, tool_name = parsed
        # mcp_navi_web_search -> navi_web_search -> split into navi, web, search
        # We try matching by removing the mcp_ prefix and comparing
        expected_old = f"mcp_{server_name}_{tool_name}"
        if expected_old == name:
            return candidate_name, candidate

    # Extra fallback: legacy colon format mcp:server:tool (from old sessions)
    legacy_matches = [
        (candidate_name, candidate)
        for candidate_name, candidate in tool_map.items()
        if is_mcp_tool(candidate_name) and name.replace(":", "__") == candidate_name
    ]
    if len(legacy_matches) == 1:
        return legacy_matches[0]

    return name, None


class ToolExecutor:
    """Runs tool calls and builds ToolEvent / Message results."""

    def __init__(self, tool_registry) -> None:
        self._tools = tool_registry

    def _maybe_background(
        self,
        tc: ToolCallRequest,
        resolved_name: str,
        tool: Tool | None,
        ctx=None,
    ) -> ToolResult | None:
        """Detach the call into the TaskManager when args carry background=true.

        Returns None when the call should run inline (no flag, unknown tool,
        tool not backgroundable). Otherwise pops the flag from a copy of the
        arguments, submits the tool execution to the TaskManager and returns
        the immediate ToolResult the agent sees.
        """
        args = dict(tc.arguments or {})
        want_background = args.pop("background", None)
        if not want_background or tool is None:
            return None
        if resolved_name not in _backgroundable_tools():
            return None
        # terminal open(background=true) has its own native meaning — start a
        # persistent detached process. Never hijack it.
        if resolved_name == "terminal" and args.get("action") == "open":
            return None

        # A background run should track the command, not the tool's short
        # foreground default (terminal 20s, code_exec 30s, ssh 60s) — otherwise
        # a "sleep 60" task is marked completed after 20s with a half-read
        # output while the process is still alive. When the agent did not pick
        # a timeout, lift it to the tools' shared maximum (300s) for the
        # detached run only.
        if resolved_name in _TIMEOUT_TOOLS and "timeout" not in args:
            args["timeout"] = _BG_TIMEOUT

        from navi.core.tasks import args_summary, get_task_manager

        session_id = getattr(ctx, "session_id", None) or current_session_id.get()
        coro_factory = lambda bg_ctx: tool.execute(args, ctx=bg_ctx)  # noqa: E731
        # Subagents emit a dense event stream — give them a deeper ring.
        ring_size = 200 if resolved_name == "spawn_agent" else None
        job = get_task_manager().submit(
            session_id=session_id,
            tool=resolved_name,
            args=args,
            coro_factory=coro_factory,
            ctx=ctx,
            parent_tool_call_id=tc.id,
            ring_size=ring_size,
        )
        if isinstance(job, str):
            # Cap/rate-limit hit — surface as a normal failed result so the
            # agent can fall back to running the tool inline.
            return ToolResult(
                success=False,
                output=f"Cannot run in background: {job}. Run it in the foreground instead.",
                error="task_cap",
            )
        output = json.dumps(
            {
                "task_id": job.task_id,
                "tool": resolved_name,
                "status": "running",
                "args_summary": args_summary(args),
                "hint": (
                    "The call is running in the background — continue with other work. "
                    "Use tasks check/wait with this task_id to collect the result; "
                    "it will also arrive as an automatic [Background task results] "
                    "note next turn."
                ),
            },
            ensure_ascii=False,
        )
        return ToolResult(success=True, output=output,
                          metadata={"background": True, "task_id": job.task_id})

    async def _execute_one(
        self,
        tc: ToolCallRequest,
        tool_map: dict[str, Tool],
        ctx=None,
    ) -> tuple["ToolEvent", Message, "Message | None"]:
        """Execute a single tool call and return (ToolEvent, tool_msg, optional_image_msg).

        This is the single canonical path for tool resolution, middleware,
        execution, and message construction. All public batch methods delegate here.
        """
        from navi.core.events import ToolEvent

        resolved_name, tool = _resolve_tool(tool_map, tc.name)
        image_msg = None
        metadata: dict = {}
        if tool is None:
            content = f"Error: tool '{tc.name}' not found."
            event = ToolEvent(
                tool_name=tc.name,
                arguments=tc.arguments,
                result=content,
                success=False,
                tool_call_id=tc.id,
            )
        else:
            background_result = self._maybe_background(tc, resolved_name, tool, ctx)
            middlewares = getattr(self._tools, "_middlewares", [])
            if background_result is not None:
                log.info(
                    "tool.backgrounded", tool=resolved_name,
                    args=tc.arguments, task_id=background_result.metadata.get("task_id"),
                )
                result = background_result
                for mw in middlewares:
                    await mw.after_execute(resolved_name, tc.arguments, result)
            else:
                log.info("tool.execute", tool=resolved_name, requested_tool=tc.name, args=tc.arguments)
                for mw in middlewares:
                    await mw.before_execute(resolved_name, tc.arguments)
                result = await tool.execute(tc.arguments, ctx=ctx)
                for mw in middlewares:
                    await mw.after_execute(resolved_name, tc.arguments, result)
            content = result.to_message_content()
            metadata = result.metadata or {}
            event = ToolEvent(
                tool_name=resolved_name,
                arguments=tc.arguments,
                result=content,
                success=result.success,
                metadata=metadata,
                tool_call_id=tc.id,
            )
            if result.success and result.metadata and result.metadata.get("is_image"):
                b64 = result.metadata.get("base64")
                if b64:
                    image_msg = Message(
                        role="user",
                        content=f"[Image loaded via {resolved_name} — analyse it]",
                        images=[b64],
                    )
        msg = Message(
            role="tool",
            content=content,
            tool_call_id=tc.id,
            name=resolved_name if tool is not None else tc.name,
            metadata=metadata,
        )
        return event, msg, image_msg

    async def _run_single_tool(
        self,
        tc: ToolCallRequest,
        tool_map: dict[str, Tool],
        ctx=None,
    ) -> tuple["ToolEvent", Message, "Message | None"]:
        """Execute one tool call and return (ToolEvent, tool_msg, optional_image_msg).

        Called via asyncio.create_task() from run_stream() so that the parent
        generator can drain the event sink queue concurrently.
        """
        return await self._execute_one(tc, tool_map, ctx=ctx)

    async def _execute_tool_calls(
        self, tool_calls: list[ToolCallRequest], tools: list[Tool], ctx=None
    ) -> tuple[list[Message], list[Message]]:
        tool_map = {t.name: t for t in tools}
        pairs = await asyncio.gather(
            *[self._execute_one(tc, tool_map, ctx=ctx) for tc in tool_calls]
        )
        tool_msgs = [p[1] for p in pairs]
        image_msgs = [p[2] for p in pairs if p[2] is not None]
        return tool_msgs, image_msgs

    async def _execute_tool_calls_streaming(
        self, tool_calls: list[ToolCallRequest], tools: list[Tool], ctx=None
    ) -> tuple[list[tuple["ToolEvent", Message]], list[Message]]:
        tool_map = {t.name: t for t in tools}
        triples = await asyncio.gather(
            *[self._execute_one(tc, tool_map, ctx=ctx) for tc in tool_calls]
        )
        return [(t[0], t[1]) for t in triples], [t[2] for t in triples if t[2] is not None]