"""MCP fleet status tool — returns all known VMs with their job associations.

Provides the ``vm_fleet_status`` tool (R005) that agents use to query the
fleet-wide VM ↔ job view.  Each VM entry includes its running state and a
list of associated jobs with their current status and last step.
"""

from __future__ import annotations

import asyncio
from typing import Any

from fastmcp.tools import ToolResult, tool
from prefab_ui.app import PrefabApp
from prefab_ui.components import Column, DataTable, DataTableColumn, ExpandableRow, Heading

from mcp_vm_blackbox.job_store import default_job_store
from mcp_vm_blackbox.models import FleetEntry, JobSummary
from mcp_vm_blackbox.vm_registry import default_registry

_MAX_STEP_LEN: int = 60


def _active_or_latest_job(jobs: list[JobSummary]) -> JobSummary | None:
    """Return the active job if present, else the most recently updated job.

    Args:
        jobs: All jobs for one VM, in any order.

    Returns:
        The single most relevant ``JobSummary`` for table display, or ``None``
        when the VM has no jobs.
    """
    if not jobs:
        return None
    for job in jobs:
        if job.status == "active":
            return job
    return max(jobs, key=lambda j: j.updated_at)


def _format_job_columns(job: JobSummary | None) -> dict[str, str]:
    """Format the job-derived display columns shared by text and table builders.

    Args:
        job: The selected ``JobSummary`` for a VM, or ``None`` when no jobs exist.

    Returns:
        A dict with keys ``job_status``, ``scenario``, ``last_step``, ``context``,
        ``updated``; all values are strings (empty string when job is ``None``).
    """
    if job is None:
        return {"job_id": "", "job_status": "", "scenario": "", "last_step": "", "context": "", "updated": ""}
    raw_step = job.last_step or ""
    last_step = (raw_step[:_MAX_STEP_LEN] + "...") if len(raw_step) > _MAX_STEP_LEN else raw_step
    context = "N/A" if job.pilot_context_usage is None else f"{int(job.pilot_context_usage * 100)}%"
    return {
        "job_id": job.job_id,
        "job_status": str(job.status),
        "scenario": job.scenario_name or "",
        "last_step": last_step,
        "context": context,
        "updated": job.updated_at or "",
    }


def _build_fleet_text(entries: list[FleetEntry]) -> str:
    """Build a plain-text summary of fleet state for agent consumption.

    Format: prose summary line + blank line + markdown table.  Empty fleet
    returns the literal string ``"No VMs registered."``.

    Args:
        entries: The assembled list of ``FleetEntry`` models.

    Returns:
        A markdown string suitable for agent parsing.
    """
    if not entries:
        return "No VMs registered."

    n_vms = len(entries)
    n_running = sum(1 for e in entries if e.running)
    n_active = sum(1 for e in entries if any(j.status == "active" for j in e.jobs))

    header = f"Fleet: {n_vms} VMs \u2014 {n_running} running, {n_vms - n_running} stopped. {n_active} active jobs."
    table_header = "| VM Name | State | Recording | Job Status | Scenario | Last Step | Context | Updated |"
    table_sep = "|---------|-------|-----------|------------|----------|-----------|---------|---------|"

    rows: list[str] = []
    for entry in entries:
        cols = _format_job_columns(_active_or_latest_job(entry.jobs))
        state = "Running" if entry.running else "Stopped"
        recording = "Yes" if entry.recording else "No"
        rows.append(
            f"| {entry.vm_name} | {state} | {recording}"
            f" | {cols['job_status']} | {cols['scenario']} | {cols['last_step']}"
            f" | {cols['context']} | {cols['updated']} |"
        )

    return "\n".join([header, "", table_header, table_sep, *rows])


def _build_fleet_table(entries: list[FleetEntry]) -> PrefabApp:
    """Build the interactive DataTable as a PrefabApp for Prefab-aware MCP clients.

    Produces one row per VM using ``_active_or_latest_job`` to flatten the
    jobs list.  VMs with no jobs produce a row with empty job columns.

    Args:
        entries: The assembled list of ``FleetEntry`` models.

    Returns:
        A ``PrefabApp`` containing a ``Heading`` and a sortable, filterable
        ``DataTable`` component.  Empty fleet produces a table with zero rows.
    """
    columns = [
        DataTableColumn(key="vm_name", header="VM Name", sortable=True),
        DataTableColumn(key="running", header="State", sortable=True),
        DataTableColumn(key="job_id", header="Active Job ID", sortable=True),
        DataTableColumn(key="recording", header="Recording", sortable=True),
        DataTableColumn(key="job_status", header="Job Status", sortable=True),
        DataTableColumn(key="scenario_name", header="Scenario", sortable=True),
        DataTableColumn(key="last_step", header="Last Step", sortable=True),
        DataTableColumn(key="context_usage", header="Context", sortable=True),
        DataTableColumn(key="updated_at", header="Updated", sortable=True),
    ]

    rows: list[dict[str, Any] | ExpandableRow] = []
    for entry in entries:
        cols = _format_job_columns(_active_or_latest_job(entry.jobs))
        rows.append({
            "vm_name": entry.vm_name,
            "running": "Running" if entry.running else "Stopped",
            "job_id": cols["job_id"],
            "recording": "Yes" if entry.recording else "No",
            "job_status": cols["job_status"],
            "scenario_name": cols["scenario"],
            "last_step": cols["last_step"],
            "context_usage": cols["context"],
            "updated_at": cols["updated"],
        })

    return PrefabApp(
        title="VM Fleet Status",
        view=Column(
            children=[
                Heading("VM Fleet Status"),
                DataTable(columns=columns, rows=rows, search=True, paginated=True, page_size=20),
            ]
        ),
    )


@tool(meta={"ui": True})
async def vm_fleet_status() -> ToolResult:
    """Return all known VMs with their running state and associated jobs.

    Reads the VM registry for the list of known VMs, then queries the job
    store for each VM to build a fleet-wide status view.

    Returns:
        A ``ToolResult`` with:
        - ``content``: Prose summary + markdown table for agent consumers
          and non-Prefab MCP clients.
        - ``structured_content``: Interactive sortable ``DataTable`` rendered
          in Prefab-aware MCP clients (e.g., Claude Code conversation UI).
          Empty fleet produces an empty table, never an error.
    """
    vm_entries = await default_registry.list_all()
    fleet: list[FleetEntry] = []
    if vm_entries:
        all_job_records = await asyncio.gather(*[default_job_store.list_by_vm(vm.vm_name) for vm in vm_entries])
        for vm, job_records in zip(vm_entries, all_job_records, strict=False):
            job_summaries = [
                JobSummary(
                    job_id=r.job_id,
                    scenario_name=r.scenario_name,
                    status=r.status,
                    last_agent_id=r.last_agent_id,
                    last_step=r.steps[-1].action if r.steps else None,
                    updated_at=r.updated_at,
                    pilot_context_usage=r.pilot_context_usage,
                )
                for r in job_records
            ]
            fleet.append(
                FleetEntry(
                    vm_name=vm.vm_name,
                    uuid=vm.uuid,
                    running=vm.running,
                    recording=vm.recording,
                    recording_path=vm.recording_path,
                    recording_started_at=vm.recording_started_at,
                    jobs=job_summaries,
                )
            )

    return ToolResult(content=_build_fleet_text(fleet), structured_content=_build_fleet_table(fleet))
