"""Cluster lifecycle tools — cluster_define, cluster_start, cluster_stop, etc.

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

from __future__ import annotations

import asyncio
import ipaddress
import logging
import os
import shutil
import subprocess
import time

import asyncssh

from vmcluster_mcp.cluster.cloud_init import generate_cloud_init_iso
from vmcluster_mcp.cluster.domain_builder import build_domain_xml
from vmcluster_mcp.cluster.libvirt_client import LibvirtClient, domain_name, generate_mac
from vmcluster_mcp.cluster.network_builder import build_network_xml, get_gateway_ip
from vmcluster_mcp.models import ClusterHandle, ClusterStatus, ErrorDetail, NodeDescriptor, NodeStatus, ToolResult
from vmcluster_mcp.server import config, mcp, structured_tool_handler
from vmcluster_mcp.topology.parser import load_topology

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

# Single shared client — connection is opened lazily on first use
_client = LibvirtClient(config.libvirt_uri)


@mcp.tool()
@structured_tool_handler
async def cluster_define(topology_name: str) -> ToolResult[ClusterStatus]:
    """Define (provision) a cluster from a topology spec.

    Creates:
    - A libvirt virtual network from the topology NetworkSpec
    - Per-node qcow2 overlay disks backed by the base image
    - Per-node cloud-init NoCloud ISOs with static IP configuration
    - libvirt domain definitions for all nodes

    Idempotent (CLUS-07): if the cluster is already defined, returns success
    with current status without modifying any resources.

    Args:
        topology_name: Name of the topology YAML in config.topology_dir.

    Returns:
        ToolResult with ClusterStatus showing all nodes as 'shut off' (newly defined)
        or current status (if idempotent path).
    """
    topo = load_topology(topology_name, config.topology_dir)

    # Idempotency check (CLUS-07): if the first domain already exists, the cluster
    # has been defined. Return current status without modifying anything.
    first_dom_name = domain_name(topo.cluster_name, topo.nodes[0].name)
    existing_dom = await _client.lookup_domain(first_dom_name)
    if existing_dom is not None:
        logger.info("cluster_define: %s already defined — returning current status", topo.cluster_name)
        net = await _client.lookup_network(topo.network.name)
        network_active = False
        if net is not None:
            network_active = await _client.is_network_active(net)
        node_statuses = []
        for node in topo.nodes:
            dom = await _client.lookup_domain(domain_name(topo.cluster_name, node.name))
            if dom is not None:
                state = await _client.get_domain_state(dom)
            else:
                state = "unknown"
            node_statuses.append(
                NodeStatus(
                    name=node.name,
                    role=node.role,
                    ip=node.ip,
                    domain_state=state,
                    ssh_reachable=False,
                )
            )
        return ToolResult(
            success=True,
            result=ClusterStatus(
                cluster_name=topo.cluster_name,
                network_active=network_active,
                nodes=node_statuses,
            ),
        )

    # Create overlay directory
    overlay_dir = os.path.join(config.overlay_dir, topo.cluster_name)
    os.makedirs(overlay_dir, exist_ok=True)

    # Define and activate the cluster network
    net_xml = build_network_xml(topo.network)
    net = await _client.define_network(net_xml)
    await _client.start_network(net)

    # Compute network parameters for cloud-init
    net_obj = ipaddress.IPv4Network(topo.network.subnet, strict=False)
    gateway = get_gateway_ip(topo.network.subnet)
    prefix_len = net_obj.prefixlen

    node_statuses = []
    errors = []

    for node in topo.nodes:
        try:
            mac = generate_mac(topo.cluster_name, node.name)
            overlay_path = os.path.join(overlay_dir, f"{node.name}.qcow2")

            # Safety assertion: base image must not equal overlay path (D-10)
            assert os.path.abspath(topo.base_image) != os.path.abspath(overlay_path), (
                f"base_image and overlay_path must differ: {topo.base_image}"
            )

            # Create backing overlay disk via qemu-img (D-07)
            subprocess.run(
                ["qemu-img", "create", "-b", topo.base_image, "-F", "qcow2", "-f", "qcow2", overlay_path],
                capture_output=True,
                check=True,
            )

            # Generate cloud-init NoCloud ISO with static IP config (D-03)
            cidata_path = generate_cloud_init_iso(
                cluster_name=topo.cluster_name,
                node_name=node.name,
                ip=node.ip,
                mac_address=mac,
                gateway_ip=gateway,
                prefix_len=prefix_len,
                output_dir=overlay_dir,
            )

            # Build and register domain XML
            dom_xml = build_domain_xml(
                topo.cluster_name,
                node,
                overlay_path,
                cidata_path,
                topo.network.bridge,
                mac,
            )
            await _client.define_domain(dom_xml)

            node_statuses.append(
                NodeStatus(
                    name=node.name,
                    role=node.role,
                    ip=node.ip,
                    domain_state="shut off",
                    ssh_reachable=False,
                )
            )
            logger.info("cluster_define: defined node %s/%s", topo.cluster_name, node.name)

        except Exception as exc:
            # D-05: partial failure — collect errors, continue to remaining nodes
            logger.error("cluster_define: failed to define node %s/%s: %s", topo.cluster_name, node.name, exc)
            errors.append(f"{node.name}: {exc}")
            node_statuses.append(
                NodeStatus(
                    name=node.name,
                    role=node.role,
                    ip=node.ip,
                    domain_state="unknown",
                    ssh_reachable=False,
                )
            )

    if errors:
        return ToolResult(
            success=False,
            result=ClusterStatus(
                cluster_name=topo.cluster_name,
                network_active=True,
                nodes=node_statuses,
            ),
            error=ErrorDetail(
                code="PARTIAL_DEFINE",
                message=f"Some nodes failed to define: {'; '.join(errors)}",
                recoverable=True,
            ),
        )

    return ToolResult(
        success=True,
        result=ClusterStatus(
            cluster_name=topo.cluster_name,
            network_active=True,
            nodes=node_statuses,
        ),
    )


# ---------------------------------------------------------------------------
# SSH readiness helper (D-04, D-11)
# ---------------------------------------------------------------------------


async def _wait_for_ssh(ip: str, port: int, user: str, key_path: str, timeout_s: int) -> tuple[bool, float]:
    """Poll SSH until the node responds or the timeout expires.

    Returns (reachable, elapsed_seconds).
    Never raises — failures are caught and retried until deadline.
    """
    start = time.monotonic()
    deadline = start + timeout_s
    while time.monotonic() < deadline:
        try:
            async with asyncssh.connect(
                ip,
                port=port,
                username=user,
                client_keys=[key_path],
                known_hosts=None,
                login_timeout=5,
            ) as conn:
                await conn.run("echo ok", check=True)
                return True, time.monotonic() - start
        except (OSError, asyncssh.Error):
            await asyncio.sleep(2)
    return False, time.monotonic() - start


# ---------------------------------------------------------------------------
# cluster_start
# ---------------------------------------------------------------------------


@mcp.tool()
@structured_tool_handler
async def cluster_start(
    cluster_name: str,
    wait_for_ssh: bool = True,
    ssh_timeout_s: int = 120,
) -> ToolResult[ClusterStatus]:
    """Start all nodes in a defined cluster.

    Boots every domain that is not already running. If wait_for_ssh is True,
    polls each node via SSH until it responds or ssh_timeout_s is exceeded.
    Per D-04: if *any* node fails SSH, returns success=False with SSH_TIMEOUT.

    Args:
        cluster_name: Name of the cluster (must already be defined).
        wait_for_ssh: Whether to wait for SSH to be reachable on all nodes.
        ssh_timeout_s: Seconds to wait per-node for SSH readiness.

    Returns:
        ToolResult with ClusterStatus; ssh_reachable set per node when wait_for_ssh=True.
    """
    topo = load_topology(cluster_name, config.topology_dir)

    # Ensure the network is active
    net = await _client.lookup_network(topo.network.name)
    network_active = False
    if net is not None:
        network_active = await _client.is_network_active(net)

    # Boot all stopped domains
    for node in topo.nodes:
        dname = domain_name(topo.cluster_name, node.name)
        dom = await _client.lookup_domain(dname)
        if dom is not None:
            active = await _client.is_domain_active(dom)
            if not active:
                await _client.start_domain(dom)
                logger.info("cluster_start: started domain %s", dname)

    # Gather per-node SSH results concurrently (D-04 / D-11)
    ssh_results: dict[str, bool] = {}
    if wait_for_ssh:
        tasks = {
            node.name: _wait_for_ssh(
                node.ip,
                22,
                topo.ssh.user,
                topo.ssh.key_path,
                ssh_timeout_s,
            )
            for node in topo.nodes
        }
        gathered = await asyncio.gather(*tasks.values())
        for node, (reachable, _) in zip(tasks.keys(), gathered):
            ssh_results[node] = reachable

    # Build node statuses
    node_statuses = []
    for node in topo.nodes:
        dname = domain_name(topo.cluster_name, node.name)
        dom = await _client.lookup_domain(dname)
        state = await _client.get_domain_state(dom) if dom is not None else "unknown"
        node_statuses.append(
            NodeStatus(
                name=node.name,
                role=node.role,
                ip=node.ip,
                domain_state=state,
                ssh_reachable=ssh_results.get(node.name, False),
            )
        )

    # D-04: strict all-or-nothing on SSH
    if wait_for_ssh and not all(ssh_results.values()):
        failed = [n for n, ok in ssh_results.items() if not ok]
        return ToolResult(
            success=False,
            result=ClusterStatus(
                cluster_name=topo.cluster_name,
                network_active=network_active,
                nodes=node_statuses,
            ),
            error=ErrorDetail(
                code="SSH_TIMEOUT",
                message=f"SSH timed out for nodes: {', '.join(failed)}",
                recoverable=True,
            ),
        )

    return ToolResult(
        success=True,
        result=ClusterStatus(
            cluster_name=topo.cluster_name,
            network_active=network_active,
            nodes=node_statuses,
        ),
    )


# ---------------------------------------------------------------------------
# cluster_stop
# ---------------------------------------------------------------------------


@mcp.tool()
@structured_tool_handler
async def cluster_stop(
    cluster_name: str,
    mode: str = "shutdown",
) -> ToolResult[ClusterStatus]:
    """Stop all running nodes in a cluster.

    Args:
        cluster_name: Name of the cluster.
        mode: 'shutdown' for ACPI graceful shutdown; 'destroy' for force-off.

    Returns:
        ToolResult with ClusterStatus showing all nodes as stopped.
    """
    topo = load_topology(cluster_name, config.topology_dir)

    net = await _client.lookup_network(topo.network.name)
    network_active = False
    if net is not None:
        network_active = await _client.is_network_active(net)

    node_statuses = []
    for node in topo.nodes:
        dname = domain_name(topo.cluster_name, node.name)
        dom = await _client.lookup_domain(dname)
        if dom is not None:
            active = await _client.is_domain_active(dom)
            if active:
                if mode == "destroy":
                    await _client.destroy_domain(dom)
                else:
                    await _client.shutdown_domain(dom)
                logger.info("cluster_stop(%s): stopped domain %s", mode, dname)
            state = await _client.get_domain_state(dom)
        else:
            state = "unknown"
        node_statuses.append(
            NodeStatus(
                name=node.name,
                role=node.role,
                ip=node.ip,
                domain_state=state,
                ssh_reachable=False,
            )
        )

    return ToolResult(
        success=True,
        result=ClusterStatus(
            cluster_name=topo.cluster_name,
            network_active=network_active,
            nodes=node_statuses,
        ),
    )


# ---------------------------------------------------------------------------
# cluster_destroy
# ---------------------------------------------------------------------------


@mcp.tool()
@structured_tool_handler
async def cluster_destroy(
    cluster_name: str,
    remove_overlays: bool = False,
) -> ToolResult[ClusterStatus]:
    """Destroy a cluster: force-stop active domains, undefine all resources.

    The network is always destroyed and undefined (D-01). Overlay disks are
    removed only when remove_overlays=True (D-02).

    Args:
        cluster_name: Name of the cluster to tear down.
        remove_overlays: If True, delete the overlay disk directory.

    Returns:
        ToolResult with ClusterStatus (all nodes removed; empty list if successful).
    """
    topo = load_topology(cluster_name, config.topology_dir)

    node_statuses = []
    for node in topo.nodes:
        dname = domain_name(topo.cluster_name, node.name)
        dom = await _client.lookup_domain(dname)
        if dom is not None:
            active = await _client.is_domain_active(dom)
            if active:
                await _client.destroy_domain(dom)
                logger.info("cluster_destroy: force-stopped domain %s", dname)
            await _client.undefine_domain(dom)
            logger.info("cluster_destroy: undefined domain %s", dname)
        node_statuses.append(
            NodeStatus(
                name=node.name,
                role=node.role,
                ip=node.ip,
                domain_state="unknown",
                ssh_reachable=False,
            )
        )

    # Always tear down the network (D-01)
    net = await _client.lookup_network(topo.network.name)
    if net is not None:
        net_active = await _client.is_network_active(net)
        if net_active:
            await _client.destroy_network(net)
        await _client.undefine_network(net)
        logger.info("cluster_destroy: removed network %s", topo.network.name)

    # Optionally remove overlay directory (D-02)
    if remove_overlays:
        overlay_dir = os.path.join(config.overlay_dir, topo.cluster_name)
        shutil.rmtree(overlay_dir, ignore_errors=True)
        logger.info("cluster_destroy: removed overlays at %s", overlay_dir)

    return ToolResult(
        success=True,
        result=ClusterStatus(
            cluster_name=topo.cluster_name,
            network_active=False,
            nodes=node_statuses,
        ),
    )


# ---------------------------------------------------------------------------
# cluster_status
# ---------------------------------------------------------------------------


@mcp.tool()
@structured_tool_handler
async def cluster_status(cluster_name: str) -> ToolResult[ClusterStatus]:
    """Return current domain state and SSH reachability for all cluster nodes.

    SSH is only checked for nodes whose domain_state is 'running'; checks run
    in parallel via asyncio.gather.

    Args:
        cluster_name: Name of the cluster to inspect.

    Returns:
        ToolResult with ClusterStatus containing per-node NodeStatus entries.
    """
    topo = load_topology(cluster_name, config.topology_dir)

    net = await _client.lookup_network(topo.network.name)
    network_active = bool(net and await _client.is_network_active(net))

    # Collect domain states for all nodes
    dom_states: list[tuple[str, str]] = []  # (node_name, domain_state)
    for node in topo.nodes:
        dname = domain_name(topo.cluster_name, node.name)
        dom = await _client.lookup_domain(dname)
        if dom is None:
            dom_states.append((node.name, "unknown"))
        else:
            state = await _client.get_domain_state(dom)
            dom_states.append((node.name, state))

    # Run SSH checks in parallel for running nodes only
    running_nodes = [
        node for node, (_, state) in zip(topo.nodes, dom_states) if state == "running"
    ]
    if running_nodes:
        ssh_tasks = [
            _wait_for_ssh(node.ip, 22, topo.ssh.user, topo.ssh.key_path, 5)
            for node in running_nodes
        ]
        ssh_gathered = await asyncio.gather(*ssh_tasks)
        ssh_reachable_map = {
            node.name: reachable for node, (reachable, _) in zip(running_nodes, ssh_gathered)
        }
    else:
        ssh_reachable_map = {}

    node_statuses = []
    for node, (_, state) in zip(topo.nodes, dom_states):
        node_statuses.append(
            NodeStatus(
                name=node.name,
                role=node.role,
                ip=node.ip,
                domain_state=state,
                ssh_reachable=ssh_reachable_map.get(node.name, False),
            )
        )

    return ToolResult(
        success=True,
        result=ClusterStatus(
            cluster_name=topo.cluster_name,
            network_active=network_active,
            nodes=node_statuses,
        ),
    )


# ---------------------------------------------------------------------------
# cluster_handle
# ---------------------------------------------------------------------------


@mcp.tool()
@structured_tool_handler
async def cluster_handle(cluster_name: str) -> ToolResult[ClusterHandle]:
    """Return a ClusterHandle descriptor for downstream SSH consumers.

    Requires the cluster to be running. Fetches the kernel version from the
    first node via SSH ('uname -r').

    Args:
        cluster_name: Name of the running cluster.

    Returns:
        ToolResult with ClusterHandle including node SSH descriptors,
        artifact_path, and kernel_version.
    """
    topo = load_topology(cluster_name, config.topology_dir)

    # Verify cluster is running — check first node
    first_dom = await _client.lookup_domain(domain_name(topo.cluster_name, topo.nodes[0].name))
    if first_dom is None or not await _client.is_domain_active(first_dom):
        return ToolResult(
            success=False,
            error=ErrorDetail(
                code="CLUSTER_NOT_RUNNING",
                message="Cluster must be running to get handle",
                recoverable=True,
            ),
        )

    # Fetch kernel version from first node
    try:
        async with asyncssh.connect(
            topo.nodes[0].ip,
            port=22,
            username=topo.ssh.user,
            client_keys=[topo.ssh.key_path],
            known_hosts=None,
            login_timeout=10,
        ) as conn:
            run_result = await conn.run("uname -r", check=True)
            kernel_version = run_result.stdout.strip()
    except (OSError, asyncssh.Error) as exc:
        return ToolResult(
            success=False,
            error=ErrorDetail(
                code="SSH_ERROR",
                message=f"Cannot determine kernel version: {exc}",
                recoverable=True,
            ),
        )

    nodes = [
        NodeDescriptor(
            name=node.name,
            role=node.role,
            ip=node.ip,
            ssh_port=22,
            ssh_user=topo.ssh.user,
            ssh_key_path=topo.ssh.key_path,
        )
        for node in topo.nodes
    ]

    return ToolResult(
        success=True,
        result=ClusterHandle(
            cluster_name=topo.cluster_name,
            nodes=nodes,
            artifact_path="/opt/vmcluster/artifacts",
            kernel_version=kernel_version,
        ),
    )


# ---------------------------------------------------------------------------
# node_crash
# ---------------------------------------------------------------------------


@mcp.tool()
@structured_tool_handler
async def node_crash(
    cluster_name: str,
    node: str,
    restart_after: bool = True,
    wait_for_ssh: bool = True,
    ssh_timeout_s: int = 120,
) -> ToolResult[NodeStatus]:
    """Simulate an unclean node failure via libvirt destroy, optionally restarting.

    Executes ``virsh destroy`` (force-off) on the target domain to simulate an
    unclean power failure. Invalidates the SSH connection cache so the next exec
    call re-establishes a fresh connection.

    Args:
        cluster_name: Name of the cluster containing the target node.
        node: Name of the node to crash.
        restart_after: If True (default), restart the node after destroying it.
        wait_for_ssh: If True and restart_after=True, poll SSH until reachable.
        ssh_timeout_s: Seconds to wait for SSH readiness after restart.

    Returns:
        ToolResult[NodeStatus] with the node's state after the operation.
        Returns success=False with SSH_TIMEOUT if the node does not become
        SSH-reachable within ssh_timeout_s.
    """
    # Lazy import: avoids circular import at module load time (exec imports server,
    # server imports cluster; importing exec from cluster at module level creates a cycle).
    # By the time node_crash is called, all modules are fully initialized.
    import libvirt  # noqa: I001
    import vmcluster_mcp.exec as _exec_pkg

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

    dname = domain_name(cluster_name, node)
    dom = await _client.lookup_domain(dname)

    # Force-off: best-effort if already shut off (D-02, D-15)
    if dom is not None:
        try:
            await _client.destroy_domain(dom)
        except libvirt.libvirtError:
            logger.warning(
                "node_crash: domain %s already shut off or destroy failed; continuing", dname
            )

    # Invalidate SSH cache so next exec re-establishes a fresh connection (D-12)
    await _exec_pkg._ssh_manager.invalidate(node_spec.ip)

    if not restart_after:
        return ToolResult(
            success=True,
            result=NodeStatus(
                name=node,
                role=node_spec.role,
                ip=node_spec.ip,
                domain_state="shut off",
                ssh_reachable=False,
            ),
        )

    # Boot the node back up (D-15)
    await _client.start_domain_by_name(dname)

    reachable = False
    if wait_for_ssh:
        reachable, _ = await _wait_for_ssh(
            node_spec.ip, 22, topo.ssh.user, topo.ssh.key_path, ssh_timeout_s
        )
        if not reachable:
            return ToolResult(
                success=False,
                error=ErrorDetail(
                    code="SSH_TIMEOUT",
                    message=f"Node {node} not SSH-reachable after {ssh_timeout_s}s",
                    recoverable=True,
                ),
            )

    # Fetch current domain state after restart
    dom_new = await _client.lookup_domain(dname)
    state = await _client.get_domain_state(dom_new) if dom_new is not None else "unknown"

    logger.info(
        "node_crash: cluster=%s node=%s restarted=True ssh_reachable=%s domain_state=%s",
        cluster_name, node, reachable, state,
    )

    return ToolResult(
        success=True,
        result=NodeStatus(
            name=node,
            role=node_spec.role,
            ip=node_spec.ip,
            domain_state=state,
            ssh_reachable=reachable,
        ),
    )
