"""Kernel observability tools — dmesg_mark, dmesg_collect.

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

from __future__ import annotations

import asyncio
import logging
import uuid

from vmcluster_mcp.exec.ssh import SSHConnectionManager
from vmcluster_mcp.models import DmesgResult, ErrorDetail, ToolResult
from vmcluster_mcp.observe.dmesg import collect_dmesg, insert_marker
from vmcluster_mcp.server import config, mcp, structured_tool_handler
from vmcluster_mcp.topology.parser import load_topology

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

_ssh_manager = SSHConnectionManager()


@mcp.tool()
@structured_tool_handler
async def dmesg_mark(
    cluster_name: str,
    nodes: list[str] | None = None,
) -> ToolResult[dict[str, str]]:
    """Insert a named dmesg marker on cluster nodes via /dev/kmsg.

    Writes a vmcluster-marker-{uuid} string to /dev/kmsg on each target node.
    All nodes in a single call receive the same marker string (single UUID per call).
    Use the returned marker strings as anchors for subsequent dmesg_collect calls.

    Args:
        cluster_name: Name of the topology (cluster) to target.
        nodes: Optional list of node names to target. If None or empty, targets all nodes.

    Returns:
        ToolResult[dict[str, str]] mapping node name to marker string for each
        successfully marked node. success=False only if ALL nodes fail.
    """
    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)

    # Generate ONE marker string for all nodes in this call (D-03)
    marker = f"vmcluster-marker-{uuid.uuid4()}"

    # Fan-out: insert marker on all target nodes concurrently (best-effort D-17)
    per_node_results: dict[str, str | Exception] = {}

    async def _mark_node(node) -> tuple[str, str | Exception]:
        try:
            await insert_marker(
                _ssh_manager,
                ip=node.ip,
                port=22,
                user=topo.ssh.user,
                key_path=topo.ssh.key_path,
                marker=marker,
            )
            return node.name, marker
        except Exception as exc:  # noqa: BLE001
            logger.warning("dmesg_mark failed on node %s: %s", node.name, exc)
            return node.name, exc

    results = await asyncio.gather(*[_mark_node(n) for n in target_nodes])

    successful: dict[str, str] = {}
    failed: list[str] = []
    for node_name, outcome in results:
        if isinstance(outcome, Exception):
            failed.append(node_name)
        else:
            successful[node_name] = outcome

    if not successful:
        return ToolResult(
            success=False,
            error=ErrorDetail(
                code="ALL_NODES_FAILED",
                message=f"dmesg_mark failed on all {len(failed)} nodes",
                recoverable=True,
            ),
        )

    return ToolResult(success=True, result=successful)


@mcp.tool()
@structured_tool_handler
async def dmesg_collect(
    cluster_name: str,
    nodes: list[str] | None = None,
    since_marker: str | None = None,
    filter_level: str = "all",
) -> ToolResult[dict[str, DmesgResult]]:
    """Collect and classify dmesg output from cluster nodes.

    Runs ``dmesg`` on each target node via SSH, parses the output, optionally
    slices results after a marker inserted by ``dmesg_mark``, and filters by
    log level.  Classification (oops/panic/kasan/ubsan/warn/normal) is applied
    to every line.

    Args:
        cluster_name: Name of the topology (cluster) to target.
        nodes: Optional list of node names. If None or empty, targets all nodes.
        since_marker: Marker string returned by a prior ``dmesg_mark`` call.
            Only lines appearing after the marker are included.
        filter_level: ``"all"`` (default) | ``"warn+"`` | ``"err+"``.

    Returns:
        ToolResult mapping node name to DmesgResult for each node.
    """
    from vmcluster_mcp.observe.dmesg import FILTER_THRESHOLDS  # local import avoids circular at module load

    if filter_level not in FILTER_THRESHOLDS:
        return ToolResult(
            success=False,
            error=ErrorDetail(
                code="INVALID_FILTER_LEVEL",
                message=f"filter_level must be one of {list(FILTER_THRESHOLDS.keys())}, got '{filter_level}'",
                recoverable=True,
            ),
        )

    topo = load_topology(cluster_name, config.topology_dir)
    topo_node_map = {n.name: n for n in topo.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)

    # Collect concurrently across nodes (D-07, D-17 best-effort)
    gather_results = await asyncio.gather(
        *[
            collect_dmesg(
                _ssh_manager,
                ip=node.ip,
                port=22,
                user=topo.ssh.user,
                key_path=topo.ssh.key_path,
                since_marker=since_marker,
                filter_level=filter_level,
            )
            for node in target_nodes
        ],
        return_exceptions=True,
    )

    results: dict[str, DmesgResult] = {}
    for node, outcome in zip(target_nodes, gather_results):
        if isinstance(outcome, Exception):
            logger.warning("dmesg_collect failed on node %s: %s", node.name, outcome)
            # Best-effort: include empty result so callers know the node was attempted
            results[node.name] = DmesgResult(
                lines=[],
                oops_detected=False,
                panic_detected=False,
                warning_count=0,
                error_count=0,
            )
        else:
            results[node.name] = outcome

    return ToolResult(success=True, result=results)
