# src/nexus_mcp/server.py
"""FastMCP server with CLI agent tools.

Exposes three MCP tools:
- batch_prompt: Send multiple prompts to CLI agents in parallel (primary tool)
- prompt: Send a single prompt to a CLI agent, routes to batch_prompt
- list_agents: Return list of supported agent names

Background task design: both prompt and batch_prompt use @mcp.tool(task=True) so they
run asynchronously and return task IDs immediately. This prevents MCP timeouts for long
operations like YOLO mode (2-5 minutes).

Testability: The tool functions are defined as plain async functions and then
registered with @mcp.tool so tests can import and call them directly without
going through the FunctionTool wrapper.
"""

import asyncio
import json
import logging
from typing import Any

from fastmcp import Context, FastMCP
from fastmcp.exceptions import ToolError
from pydantic import ValidationError

from nexus_mcp.config import get_tool_timeout
from nexus_mcp.runners.factory import RunnerFactory
from nexus_mcp.types import (
    DEFAULT_MAX_CONCURRENCY,
    AgentTask,
    AgentTaskResult,
    ExecutionMode,
    MultiPromptResponse,
    SessionPreferences,
)

mcp = FastMCP("nexus-mcp")
logger = logging.getLogger(__name__)

_PREFERENCES_KEY = "nexus:preferences"


async def _get_session_preferences(ctx: Context | None) -> SessionPreferences:
    """Read session preferences from ctx state, returning defaults when unset or ctx is None."""
    if ctx is None:
        return SessionPreferences()
    raw = await ctx.get_state(_PREFERENCES_KEY)
    if raw is None:
        return SessionPreferences()
    try:
        return SessionPreferences(**raw)  # reconstruct from dict after JSON round-trip
    except (ValidationError, TypeError) as e:
        raise ToolError(f"Session preferences are corrupted and cannot be loaded: {e}") from e


def _apply_preferences(task: AgentTask, prefs: SessionPreferences) -> AgentTask:
    """Fill in None fields on a task from session preferences.

    Returns the same task object if no updates are needed, otherwise a model_copy.
    Primary resolution happens in prompt(); this is the safety net for batch_prompt tasks.
    """
    updates: dict[str, Any] = {}
    if task.execution_mode is None:
        updates["execution_mode"] = prefs.execution_mode or "default"
    if task.model is None and prefs.model is not None:
        updates["model"] = prefs.model
    if updates:
        return task.model_copy(update=updates)
    return task


def _next_available_label(base: str, reserved: set[str]) -> str:
    """Return the first available label derived from base, avoiding reserved names.

    Returns base if available, otherwise base-2, base-3, etc.

    Args:
        base: Preferred label (typically agent name).
        reserved: Set of already-taken labels.

    Returns:
        base if not reserved, otherwise base-N for the lowest N >= 2 not in reserved.
    """
    if base not in reserved:
        return base
    n = 2
    while f"{base}-{n}" in reserved:
        n += 1
    return f"{base}-{n}"


def _assign_labels(tasks: list[AgentTask]) -> list[AgentTask]:
    """Assign unique labels to tasks, preserving explicit ones.

    Two-pass algorithm:
    1. Reserve all explicit labels
    2. Auto-assign from agent name with -N suffixes for collisions

    Args:
        tasks: List of AgentTask objects (may have label=None).

    Returns:
        New list of AgentTask objects with label set on every item.
        Input tasks are never mutated (model_copy() used throughout).
    """
    reserved = {t.label for t in tasks if t.label is not None}

    result: list[AgentTask] = []
    for task in tasks:
        if task.label is not None:
            result.append(task)
            continue

        label = _next_available_label(task.agent, reserved)
        reserved.add(label)
        result.append(task.model_copy(update={"label": label}))

    return result


async def batch_prompt(
    tasks: list[AgentTask],
    max_concurrency: int = DEFAULT_MAX_CONCURRENCY,
    ctx: Context | None = None,
) -> MultiPromptResponse:
    """Send multiple prompts to CLI agents in parallel (primary tool).

    Fans out tasks server-side with asyncio.gather and a semaphore, enabling
    true parallel agent execution within a single MCP call. Single-task usage
    is perfectly valid — use prompt for convenience when sending one task.

    Args:
        tasks: List of AgentTask objects, each with agent, prompt, and optional fields.
        max_concurrency: Max parallel agent invocations (default: 3).
        ctx: MCP context (auto-injected by FastMCP). None when called directly in tests.

    Returns:
        MultiPromptResponse with results for each task.
    """
    # Docket serializes/deserializes task arguments as JSON when task=True,
    # converting AgentTask objects to plain dicts. Reconstruct them here.
    tasks = [AgentTask(**t) if isinstance(t, dict) else t for t in tasks]

    # Resolve session preferences before dispatching — workers in a separate process
    # (Redis-backed Docket) cannot access the foreground session's MemoryStore.
    prefs = await _get_session_preferences(ctx)
    tasks = [_apply_preferences(t, prefs) for t in tasks]

    labelled = _assign_labels(tasks)
    semaphore = asyncio.Semaphore(max_concurrency)
    completed = 0

    if ctx:
        await ctx.info(f"Starting batch of {len(labelled)} tasks (concurrency={max_concurrency})")

    async def _run_single(task: AgentTask) -> AgentTaskResult:
        nonlocal completed
        async with semaphore:
            try:
                request = task.to_request()
                runner = RunnerFactory.create(task.agent)
                response = await runner.run(request)
                return AgentTaskResult(label=task.label, output=response.output)  # type: ignore[arg-type]
            except Exception as e:
                logger.exception("Task %r failed: %s", task.label, e)
                return AgentTaskResult(label=task.label, error=str(e), error_type=type(e).__name__)  # type: ignore[arg-type]
            finally:
                completed += 1
                if ctx:
                    await ctx.report_progress(
                        progress=completed,
                        total=len(labelled),
                        message=f"Completed task {completed}/{len(labelled)}: {task.label}",
                    )

    results = await asyncio.gather(*[_run_single(t) for t in labelled])
    response = MultiPromptResponse(results=list(results))
    if ctx:
        await ctx.info(f"Batch complete: {response.succeeded}/{response.total} succeeded")
    return response


async def prompt(
    agent: str,
    prompt: str,
    context: dict[str, Any] | None = None,
    execution_mode: ExecutionMode | None = None,
    model: str | None = None,
    max_retries: int | None = None,
    ctx: Context | None = None,
) -> str:
    """Send a prompt to a CLI agent as a background task.

    Returns immediately with a task ID. Client polls for results.
    This prevents timeouts for long operations (YOLO mode: 2-5 minutes).

    Args:
        agent: Agent name (e.g., "gemini")
        prompt: Prompt text to send to the agent
        context: Optional context metadata
        execution_mode: 'default' (safe), 'sandbox', or 'yolo'. None inherits session preference.
        model: Optional model name. None inherits session preference or uses CLI default.
        max_retries: Max retry attempts for transient errors (None uses env default)
        ctx: MCP context (auto-injected by FastMCP). None when called directly in tests.

    Returns:
        Agent's response text
    """
    # Resolve session preferences here (foreground) so concrete values reach the Docket worker.
    prefs = await _get_session_preferences(ctx)
    session_mode = prefs.execution_mode or "default"
    resolved_mode = execution_mode if execution_mode is not None else session_mode
    resolved_model = model if model is not None else prefs.model

    task = AgentTask(
        agent=agent,
        prompt=prompt,
        context=context or {},
        execution_mode=resolved_mode,
        model=resolved_model,
        max_retries=max_retries,
    )
    result = await batch_prompt(tasks=[task], ctx=ctx)
    task_result = result.results[0]
    if task_result.error:
        raise ToolError(task_result.formatted_error)
    return task_result.output  # type: ignore[return-value]


def list_agents() -> list[str]:
    """Return list of supported agent names.

    Returns:
        List of agent names that can be used with prompt or batch_prompt.
    """
    return RunnerFactory.list_agents()


async def set_preferences(
    execution_mode: ExecutionMode | None = None,
    model: str | None = None,
    clear_execution_mode: bool = False,
    clear_model: bool = False,
    ctx: Context | None = None,
) -> str:
    """Set session-scoped preferences that apply to subsequent prompt/batch_prompt calls.

    Preferences persist for the duration of the MCP session. Call again to update,
    or use clear_preferences to reset all fields at once.

    To clear a single field while keeping others, pass the corresponding clear_* flag:
        set_preferences(clear_model=True)  # clears model, keeps execution_mode

    Args:
        execution_mode: Default execution mode for this session ('default', 'sandbox', 'yolo').
            None retains the current session value (use clear_execution_mode=True to reset).
        model: Default model name for this session (e.g. 'gemini-2.5-flash').
            None retains the current session value (use clear_model=True to reset).
        clear_execution_mode: If True, resets execution_mode to None regardless of the
            execution_mode argument.
        clear_model: If True, resets model to None regardless of the model argument.
        ctx: MCP context (auto-injected by FastMCP).

    Returns:
        Confirmation string with the active preferences as JSON.
    """
    if ctx is None:
        raise ToolError("set_preferences requires an active session context")
    existing = await _get_session_preferences(ctx)

    new_execution_mode: ExecutionMode | None
    if clear_execution_mode:
        new_execution_mode = None
    elif execution_mode is not None:
        new_execution_mode = execution_mode
    else:
        new_execution_mode = existing.execution_mode

    new_model: str | None
    if clear_model:
        new_model = None
    elif model is not None:
        new_model = model
    else:
        new_model = existing.model

    merged = SessionPreferences(execution_mode=new_execution_mode, model=new_model)
    await ctx.set_state(_PREFERENCES_KEY, merged.model_dump())
    return f"Preferences set: {json.dumps(merged.model_dump())}"


async def get_preferences(ctx: Context | None = None) -> dict[str, Any]:
    """Return the current session preferences.

    Returns:
        Dict with 'execution_mode' and 'model' keys (None when unset).
    """
    if ctx is None:
        raise ToolError("get_preferences requires an active session context")
    prefs = await _get_session_preferences(ctx)
    return prefs.model_dump()


async def clear_preferences(ctx: Context | None = None) -> str:
    """Clear all session preferences, reverting to per-call defaults.

    Returns:
        Confirmation string.
    """
    if ctx is None:
        raise ToolError("clear_preferences requires an active session context")
    await ctx.delete_state(_PREFERENCES_KEY)
    return "Preferences cleared"


# Register functions as MCP tools after definition so tests can import
# and call the raw functions directly (not the FunctionTool wrappers).
# timeout wraps synchronous calls with anyio.fail_after(); when a client
# passes task=True, the call goes through Docket instead (subprocess-level
# timeout applies). Both prompt and batch_prompt support either path.
_tool_timeout = get_tool_timeout()
mcp.tool(task=True, timeout=_tool_timeout)(batch_prompt)
mcp.tool(task=True, timeout=_tool_timeout)(prompt)
mcp.tool()(list_agents)
mcp.tool()(set_preferences)
mcp.tool()(get_preferences)
mcp.tool()(clear_preferences)
