"""MCP tool bridge — routes agent tool calls to IDE-registered MCP servers.

The IDE frontend sends available MCP tool schemas to the backend in the chat
request's ``tools`` field.  When the LLM invokes one of these tools during the
tool-use loop, we:

1. Emit an ``mcp_invoke`` SSE event with the tool name and arguments
2. Wait for the IDE to POST the result back to /api/mcp/result/{call_id}
3. Return the result to the tool-use loop

This creates a closed feedback loop: LLM → backend → SSE → IDE MCP client →
POST result → backend → LLM.
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Awaitable

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Built-in tool names that should NOT be routed to MCP
# ---------------------------------------------------------------------------

BUILTIN_TOOL_NAMES: set[str] = {
	"read_file", "write_file", "edit_file", "run_command",
	"search_files", "list_files", "get_diagnostics", "get_file_outline",
	"run_scan", "web_search", "web_fetch", "run_tests",
	"browser_navigate", "browser_screenshot", "git_operations",
}

# ---------------------------------------------------------------------------
# Internal types
# ---------------------------------------------------------------------------


@dataclass
class _PendingInvocation:
	"""Tracks a pending MCP tool invocation awaiting IDE response."""

	event: asyncio.Event = field(default_factory=asyncio.Event)
	result: str = ""
	is_error: bool = False
	created_at: float = field(default_factory=time.monotonic)


# ---------------------------------------------------------------------------
# Module-level state
# ---------------------------------------------------------------------------

# Per-session MCP tool registries:  session_id → list of tool schemas
_session_tools: dict[str, list[dict[str, Any]]] = {}

# Pending MCP invocations awaiting IDE response:  call_id → _PendingInvocation
_pending: dict[str, _PendingInvocation] = {}

# ---------------------------------------------------------------------------
# Session tool registry
# ---------------------------------------------------------------------------


def register_session_tools(session_id: str, tools: list[dict[str, Any]]) -> None:
	"""Store MCP tool schemas from the frontend for the duration of a chat session.

	Only tools whose ``name`` is *not* in :data:`BUILTIN_TOOL_NAMES` are kept —
	built-in tools are handled directly by the tool-use loop.
	"""
	mcp_tools = [t for t in tools if t.get("name") not in BUILTIN_TOOL_NAMES]
	_session_tools[session_id] = mcp_tools
	if mcp_tools:
		names = [t.get("name", "<unnamed>") for t in mcp_tools]
		logger.info("MCP bridge: registered %d tools for session %s: %s", len(mcp_tools), session_id, names)


def get_session_tools(session_id: str) -> list[dict[str, Any]]:
	"""Return registered MCP tools for a session (empty list if none)."""
	return _session_tools.get(session_id, [])


def get_session_tool(session_id: str, tool_name: str) -> dict[str, Any] | None:
	"""Look up a single MCP tool schema by name within a session."""
	for tool in _session_tools.get(session_id, []):
		if tool.get("name") == tool_name:
			return tool
	return None


def is_mcp_tool(session_id: str, tool_name: str) -> bool:
	"""Return *True* if *tool_name* is a registered MCP tool (not a built-in)."""
	if tool_name in BUILTIN_TOOL_NAMES:
		return False
	return get_session_tool(session_id, tool_name) is not None


def register_skill_tools(
	session_id: str,
	skill_tools: list[dict[str, Any]],
) -> None:
	"""Register tools declared in active skill frontmatter (``tools_provided``).

	Each tool becomes available to the LLM via the MCP dispatch path.
	If a tool references an ``mcp_server`` + ``mcp_tool``, those are stored
	as metadata so the IDE can route accordingly.  Otherwise, the tool is
	treated as a virtual MCP tool and the IDE will execute it via the
	standard MCP invocation flow.

	Duplicate names (already registered for this session) are skipped.
	"""
	existing = {t.get("name") for t in _session_tools.get(session_id, [])}
	new_tools: list[dict[str, Any]] = []
	for st in skill_tools:
		name = st.get("name", "")
		if not name or name in BUILTIN_TOOL_NAMES or name in existing:
			continue
		tool_schema: dict[str, Any] = {
			"name": name,
			"description": st.get("description", ""),
			"input_schema": {"type": "object", "properties": {}},
			"_skill_origin": True,
		}
		if st.get("mcp_server"):
			tool_schema["_mcp_server"] = st["mcp_server"]
			tool_schema["_mcp_tool"] = st.get("mcp_tool", name)
		new_tools.append(tool_schema)
		existing.add(name)

	if new_tools:
		current = _session_tools.get(session_id, [])
		_session_tools[session_id] = current + new_tools
		names = [t["name"] for t in new_tools]
		logger.info("MCP bridge: registered %d skill tools for session %s: %s", len(new_tools), session_id, names)


def clear_session(session_id: str) -> None:
	"""Remove all registered MCP tools for a session and clean up pending invocations."""
	_session_tools.pop(session_id, None)
	# Clean up any lingering pending invocations for this session.
	# (call_ids are opaque — we cannot scope them to sessions, so we only
	# remove stale entries older than 60 s as a safety net.)
	_evict_stale_pending(max_age=60.0)
	logger.info("MCP bridge: cleared session %s", session_id)

# ---------------------------------------------------------------------------
# MCP invocation handling
# ---------------------------------------------------------------------------


async def invoke_mcp_tool(
	session_id: str,
	call_id: str,
	tool_name: str,
	args: dict[str, Any],
	sse_callback: Callable[[dict[str, Any]], Awaitable[None]],
	timeout: float = 30.0,
) -> str:
	"""Emit an ``mcp_invoke`` SSE event and block until the IDE posts the result.

	Parameters
	----------
	session_id:
		Active chat session identifier.
	call_id:
		Unique identifier for this tool invocation (generated by the caller).
	tool_name:
		Name of the MCP tool to invoke.
	args:
		Arguments to pass to the MCP tool.
	sse_callback:
		Async callable that sends an SSE event dict to the connected client.
		Expected payload shape::

			{
				"type": "mcp_invoke",
				"call_id": "...",
				"tool_name": "...",
				"arguments": { ... }
			}

	timeout:
		Maximum seconds to wait for the IDE to post the result.

	Returns
	-------
	str
		The tool execution result (text).

	Raises
	------
	TimeoutError
		If the IDE does not respond within *timeout* seconds.
	RuntimeError
		If the tool is not registered for the session.
	"""
	if not is_mcp_tool(session_id, tool_name):
		raise RuntimeError(f"Tool '{tool_name}' is not a registered MCP tool for session {session_id}")

	# Create a pending invocation entry
	pending = _PendingInvocation()
	_pending[call_id] = pending

	# Emit the SSE event so the IDE picks it up
	sse_payload: dict[str, Any] = {
		"type": "mcp_invoke",
		"call_id": call_id,
		"tool_name": tool_name,
		"arguments": args,
	}
	logger.info("MCP bridge: invoking tool '%s' (call_id=%s)", tool_name, call_id)
	await sse_callback(sse_payload)

	# Wait for the IDE to POST the result back
	try:
		await asyncio.wait_for(pending.event.wait(), timeout=timeout)
	except asyncio.TimeoutError:
		_pending.pop(call_id, None)
		logger.warning("MCP bridge: timeout waiting for result of tool '%s' (call_id=%s)", tool_name, call_id)
		raise TimeoutError(
			f"MCP tool '{tool_name}' did not return a result within {timeout}s (call_id={call_id})"
		)

	# Retrieve and clean up
	result = pending.result
	is_error = pending.is_error
	_pending.pop(call_id, None)

	if is_error:
		logger.warning("MCP bridge: tool '%s' returned error (call_id=%s): %s", tool_name, call_id, result[:200])
	else:
		logger.info("MCP bridge: tool '%s' returned successfully (call_id=%s, %d chars)", tool_name, call_id, len(result))

	return result


def submit_mcp_result(call_id: str, result: str, is_error: bool = False) -> bool:
	"""Called by the REST endpoint when the IDE posts back an MCP tool result.

	Parameters
	----------
	call_id:
		The invocation identifier matching the ``mcp_invoke`` SSE event.
	result:
		The tool execution result text.
	is_error:
		Whether the result represents an error.

	Returns
	-------
	bool
		*True* if the result was accepted (a pending invocation existed),
		*False* if no pending invocation was found for *call_id*.
	"""
	pending = _pending.get(call_id)
	if pending is None:
		logger.warning("MCP bridge: no pending invocation for call_id=%s (may have timed out)", call_id)
		return False

	pending.result = result
	pending.is_error = is_error
	pending.event.set()
	return True

# ---------------------------------------------------------------------------
# Housekeeping
# ---------------------------------------------------------------------------


def _evict_stale_pending(max_age: float = 120.0) -> None:
	"""Remove pending invocations older than *max_age* seconds."""
	now = time.monotonic()
	stale = [cid for cid, p in _pending.items() if (now - p.created_at) > max_age]
	for cid in stale:
		_pending.pop(cid, None)
	if stale:
		logger.info("MCP bridge: evicted %d stale pending invocations", len(stale))
