"""SSH execution tools — node_exec, node_exec_all.

Tools register on the FastMCP server instance imported from server.py.
All tool handlers are wrapped with @structured_tool_handler.
"""

from __future__ import annotations

import asyncio
import logging

from vmcluster_mcp.exec.ssh import SSHConnectionManager
from vmcluster_mcp.models import ErrorDetail, ExecAllResult, ExecResult, ToolResult
from vmcluster_mcp.server import config, mcp, structured_tool_handler
from vmcluster_mcp.topology.parser import load_topology

logger = logging.getLogger("vmcluster_mcp.exec")

# Module-level SSH connection manager — shared across all exec tool calls
_ssh_manager = SSHConnectionManager()


@mcp.tool()
@structured_tool_handler
async def node_exec(
    cluster_name: str,
    node_name: str,
    command: str,
    timeout_s: int = 30,
) -> ToolResult[ExecResult]:
    """Run a shell command on a single named cluster node.

    Args:
        cluster_name: Name of the topology (cluster) to target.
        node_name: Name of the node within the cluster to run the command on.
        command: Shell command string to execute on the remote node.
        timeout_s: Command execution timeout in seconds (default: 30).

    Returns:
        ToolResult[ExecResult] with exit_code, stdout, stderr, duration_s, timed_out.
        success=True even on non-zero exit codes — use exit_code to inspect command result.
        success=False only when the node cannot be found or an infrastructure error occurs.
    """
    topo = load_topology(cluster_name, config.topology_dir)

    node = next((n for n in topo.nodes if n.name == node_name), None)
    if node is None:
        return ToolResult(
            success=False,
            error=ErrorDetail(
                code="NODE_NOT_FOUND",
                message=f"Node '{node_name}' not found in cluster '{cluster_name}'",
                node=node_name,
                recoverable=False,
            ),
        )

    exec_result = await _ssh_manager.execute(
        ip=node.ip,
        port=22,
        user=topo.ssh.user,
        key_path=topo.ssh.key_path,
        command=command,
        timeout_s=timeout_s,
    )
    return ToolResult(success=True, result=exec_result)


@mcp.tool()
@structured_tool_handler
async def node_exec_all(
    cluster_name: str,
    command: str,
    nodes: list[str] | None = None,
    require_all: bool = True,
    timeout_s: int = 30,
) -> ToolResult[ExecAllResult]:
    """Run a shell command across all (or a subset of) cluster nodes in parallel.

    Args:
        cluster_name: Name of the topology (cluster) to target.
        command: Shell command string to execute on every target node.
        nodes: Optional list of node names to target. If None or empty, targets all nodes.
        require_all: When True (default), returns success=False if any node fails.
                     When False, returns success=True even with partial failures.
        timeout_s: Per-node command execution timeout in seconds (default: 30).

    Returns:
        ToolResult[ExecAllResult] with per-node ExecResult map, all_succeeded flag,
        and failed_nodes list.
    """
    topo = load_topology(cluster_name, config.topology_dir)

    topo_node_map = {n.name: n for n in topo.nodes}

    # Resolve target nodes
    if nodes:
        missing = [name for name in nodes if name not in topo_node_map]
        if missing:
            return ToolResult(
                success=False,
                error=ErrorDetail(
                    code="NODE_NOT_FOUND",
                    message=f"Nodes not found in cluster '{cluster_name}': {', '.join(missing)}",
                    recoverable=False,
                ),
            )
        target_nodes = [topo_node_map[name] for name in nodes]
    else:
        target_nodes = list(topo.nodes)

    # Fan out via asyncio.gather for concurrent execution
    tasks = [
        _ssh_manager.execute(
            ip=node.ip,
            port=22,
            user=topo.ssh.user,
            key_path=topo.ssh.key_path,
            command=command,
            timeout_s=timeout_s,
        )
        for node in target_nodes
    ]
    exec_results: list[ExecResult] = await asyncio.gather(*tasks)

    results: dict[str, ExecResult] = {
        node.name: result for node, result in zip(target_nodes, exec_results)
    }
    failed_nodes = [
        name for name, r in results.items() if r.exit_code != 0 or r.timed_out
    ]
    all_succeeded = len(failed_nodes) == 0

    if require_all and not all_succeeded:
        return ToolResult(
            success=False,
            result=ExecAllResult(results=results, all_succeeded=False, failed_nodes=failed_nodes),
            error=ErrorDetail(
                code="PARTIAL_FAILURE",
                message=f"Command failed on nodes: {', '.join(failed_nodes)}",
                recoverable=True,
            ),
        )

    return ToolResult(
        success=True,
        result=ExecAllResult(results=results, all_succeeded=all_succeeded, failed_nodes=failed_nodes),
    )
