Agentic-OS/backend/app/agent/proxmox_profiles.py

"""Proxmox connect strategy: REST API (proxmox-mcp) first, SSH fallback (ssh-generic-mcp)."""
from __future__ import annotations

import json
import re
from typing import Awaitable, Callable

from app.agent.mcp_helpers import tool_result_ok
from app.agent.tool_runner import invoke_tool
from app.services.device_defaults import proxmox_api_configured, resolve_proxmox_api_defaults

EmitFn = Callable[[int, str, str, dict | None], Awaitable[None]]

PROXMOX_MCP = "proxmox-mcp"
SSH_MCP = "ssh-generic-mcp"

MAX_VM_INVENTORY_CALLS = 40

# Node-less probe — VM/LXC calls need a resolved node name.
PROXMOX_API_DIAGNOSTICS: list[tuple[str, Callable[[dict], dict]]] = [
    ("proxmox_list_nodes", lambda _c: {}),
]

# Minimal per-node summary (health-style tasks).
PROXMOX_API_NODE_SUMMARY: list[str] = [
    "proxmox_list_qemu",
    "proxmox_list_lxc",
    "proxmox_get_node_status",
]

PROXMOX_SSH_DIAGNOSTICS: list[tuple[str, Callable[[dict], dict]]] = [
    ("ssh_run", lambda _c: {"command": "pvesh get /nodes --output-format json 2>/dev/null | head -c 4000 || true"}),
    ("ssh_run", lambda _c: {"command": "qm list 2>/dev/null || true"}),
    ("ssh_run", lambda _c: {"command": "pct list 2>/dev/null || true"}),
    ("ssh_run", lambda _c: {"command": "uptime && free -m && df -h /"}),
]

_VM_INVENTORY_TERMS = (
    "config",
    "configuration",
    "each vm",
    "every vm",
    "all vm",
    "all vms",
    "running vm",
    "running vms",
    "guest",
    "inventory",
    "vm status",
    "vm list",
    "list vm",
    "list vms",
    "qemu",
    "lxc",
)


def enrich_proxmox_device(device: dict) -> dict:
    """Attach Proxmox API env defaults; URL uses vessel public IP (:8006) unless overridden in .env."""
    public_ip = device.get("address") or ""
    api = resolve_proxmox_api_defaults(public_ip)
    return {**device, **api}


def _text_has_term(text: str, term: str) -> bool:
    term = term.lower().strip()
    if not term:
        return False
    if " " in term:
        return term in text
    if len(term) <= 2:
        return term in text.split()
    return bool(re.search(rf"\b{re.escape(term)}\b", text))


def wants_vm_inventory(title: str, issue: str) -> bool:
    """True when the user asked for per-VM status and/or config details."""
    text = f"{title} {issue}".lower()
    if any(_text_has_term(text, term) for term in _VM_INVENTORY_TERMS):
        return True
    # "VM" / "VMs" in title or issue (e.g. "Proxmox VM status").
    return bool(re.search(r"\bvms?\b", text))


def _ssh_connect_args(ctx: dict) -> dict:
    return {
        "host": ctx["address"],
        "port": ctx.get("port") or 22,
        "username": ctx.get("username") or "root",
        "password": ctx.get("secret"),
        "name": ctx.get("catalog_key") or "proxmox",
    }


def _api_connect_args(ctx: dict) -> dict:
    return {
        "url": ctx["proxmox_api_url"],
        "token_id": ctx["proxmox_token_id"],
        "token_secret": ctx["proxmox_token_secret"],
        "verify_ssl": bool(ctx.get("proxmox_verify_ssl")),
        "name": ctx.get("catalog_key") or "proxmox",
        "set_active": True,
    }


async def _server_ready(manager, server: str) -> bool:
    state = manager.servers.get(server or "")
    return bool(server and state and state.status == "loaded")


async def _try_proxmox_api(manager, device: dict) -> bool:
    if not proxmox_api_configured(device):
        return False
    if not await _server_ready(manager, PROXMOX_MCP):
        return False

    connect_res = await manager.call_tool(PROXMOX_MCP, "proxmox_connect", _api_connect_args(device))
    if not tool_result_ok(connect_res):
        return False

    probe = await manager.call_tool(PROXMOX_MCP, "proxmox_list_nodes", {})
    if tool_result_ok(probe):
        return True

    name = device.get("catalog_key") or "proxmox"
    try:
        await manager.call_tool(PROXMOX_MCP, "proxmox_disconnect_target", {"name": name})
    except Exception:  # noqa: BLE001
        pass
    return False


async def _try_proxmox_ssh(manager, device: dict) -> bool:
    if not await _server_ready(manager, SSH_MCP):
        return False
    if not device.get("secret"):
        return False

    connect_res = await manager.call_tool(SSH_MCP, "ssh_connect", _ssh_connect_args(device))
    return tool_result_ok(connect_res)


def _first_proxmox_node(output: str) -> str | None:
    text = (output or "").strip()
    if not text:
        return None
    try:
        data = json.loads(text)
        if isinstance(data, list) and data:
            item = data[0]
            if isinstance(item, dict):
                return item.get("node") or item.get("name")
        if isinstance(data, dict):
            nodes = data.get("data") or data.get("nodes")
            if isinstance(nodes, list) and nodes:
                item = nodes[0]
                if isinstance(item, dict):
                    return item.get("node") or item.get("name")
    except json.JSONDecodeError:
        pass
    match = re.search(r'"node"\s*:\s*"([^"]+)"', text)
    return match.group(1) if match else None


def _parse_qemu_vms(output: str) -> list[dict]:
    """Extract VM entries from proxmox_list_qemu JSON or `qm list` text."""
    text = (output or "").strip()
    if not text:
        return []

    parsed = None
    try:
        parsed = json.loads(text)
    except json.JSONDecodeError:
        start, end = text.find("["), text.rfind("]")
        if start != -1 and end > start:
            try:
                parsed = json.loads(text[start : end + 1])
            except json.JSONDecodeError:
                parsed = None

    rows: list[dict] = []
    if isinstance(parsed, list):
        for item in parsed:
            if isinstance(item, dict) and item.get("vmid") is not None:
                rows.append(item)
        if rows:
            return rows
    if isinstance(parsed, dict):
        data = parsed.get("data")
        if isinstance(data, list):
            for item in data:
                if isinstance(item, dict) and item.get("vmid") is not None:
                    rows.append(item)
            if rows:
                return rows

    # `qm list` tabular output
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("VMID") or line.startswith("-"):
            continue
        parts = line.split()
        if not parts or not parts[0].isdigit():
            continue
        vmid = int(parts[0])
        status = parts[2] if len(parts) > 2 else ""
        name = parts[1] if len(parts) > 1 else ""
        rows.append({"vmid": vmid, "name": name, "status": status})

    return rows


def _api_guest_inventory_empty(results: list[dict]) -> bool:
    """True when API list_qemu/list_lxc ran but returned no guests (common RBAC gap on /vms)."""
    qemu_out = next((r.get("output") for r in results if r.get("tool") == "proxmox_list_qemu"), None)
    lxc_out = next((r.get("output") for r in results if r.get("tool") == "proxmox_list_lxc"), None)
    if qemu_out is None and lxc_out is None:
        return False
    qemu_empty = not _parse_qemu_vms(qemu_out or "")
    lxc_empty = not _parse_qemu_vms(lxc_out or "")
    return qemu_empty and lxc_empty


def _vms_for_detail(vms: list[dict], title: str, issue: str) -> list[dict]:
    """Pick which guests need per-VM status/config calls."""
    text = f"{title} {issue}".lower()
    running_only = "running" in text and "all vm" not in text and "every vm" not in text
    if running_only:
        picked = [v for v in vms if str(v.get("status", "")).lower() == "running"]
        return picked or vms
    return vms


async def _append_tool_result(
    results: list[dict],
    *,
    task_id: int,
    label: str,
    server: str,
    tool: str,
    args: dict,
    manager,
    emit: EmitFn,
    mode: str,
    intent: str,
    task_title: str,
    task_issue: str,
) -> bool:
    """Run one tool and append to results. Returns False if max budget exhausted."""
    res, tool_ok = await invoke_tool(
        task_id,
        label,
        server,
        tool,
        args,
        manager,
        emit,
        task_title=task_title,
        task_issue=task_issue,
    )
    results.append(
        {
            "device": label,
            "tool": tool,
            "ok": tool_ok,
            "output": res.get("text", ""),
            "connect_mode": mode,
            "intent": intent,
            "vmid": args.get("vmid"),
        }
    )
    return True


async def _run_api_node_tools(
    results: list[dict],
    *,
    task_id: int,
    label: str,
    server: str,
    node: str,
    manager,
    emit: EmitFn,
    mode: str,
    intent: str,
    task_title: str,
    task_issue: str,
    max_diagnostics: int,
    vm_inventory: bool,
) -> None:
    node_tools = ["proxmox_list_qemu", "proxmox_list_lxc"]
    if not vm_inventory:
        node_tools.append("proxmox_get_node_status")

    qemu_output = ""
    for tool in node_tools:
        if len(results) >= max_diagnostics:
            break
        await _append_tool_result(
            results,
            task_id=task_id,
            label=label,
            server=server,
            tool=tool,
            args={"node": node},
            manager=manager,
            emit=emit,
            mode=mode,
            intent=intent,
            task_title=task_title,
            task_issue=task_issue,
        )
        if tool == "proxmox_list_qemu":
            qemu_output = results[-1].get("output") or ""

    if not vm_inventory or len(results) >= max_diagnostics:
        return

    vms = _vms_for_detail(_parse_qemu_vms(qemu_output), task_title, task_issue)
    for vm in vms[:15]:
        vmid = vm.get("vmid")
        if vmid is None or len(results) >= max_diagnostics:
            break
        for tool in ("proxmox_get_qemu_status", "proxmox_get_qemu_config"):
            if len(results) >= max_diagnostics:
                break
            await _append_tool_result(
                results,
                task_id=task_id,
                label=label,
                server=server,
                tool=tool,
                args={"node": node, "vmid": int(vmid)},
                manager=manager,
                emit=emit,
                mode=mode,
                intent=intent,
                task_title=task_title,
                task_issue=task_issue,
            )


async def _run_ssh_inventory_fallback(
    results: list[dict],
    *,
    task_id: int,
    label: str,
    device: dict,
    manager,
    emit: EmitFn,
    intent: str,
    task_title: str,
    task_issue: str,
    max_diagnostics: int,
    vm_inventory: bool,
) -> bool:
    """Use SSH qm/pct when API connected but guest lists are empty (token lacks /vms VM.Audit)."""
    ctx = enrich_proxmox_device(device)
    if not ctx.get("secret"):
        return False
    if not await _server_ready(manager, SSH_MCP):
        return False
    if not await _try_proxmox_ssh(manager, ctx):
        return False

    await emit(
        task_id,
        "connect",
        f"{label}: API guest list empty (token likely lacks /vms VM.Audit) — using SSH",
        {"device": label, "mode": "ssh_fallback"},
    )

    for command in ("qm list 2>/dev/null || true", "pct list 2>/dev/null || true"):
        if len(results) >= max_diagnostics:
            break
        await _append_tool_result(
            results,
            task_id=task_id,
            label=label,
            server=SSH_MCP,
            tool="ssh_run",
            args={"command": command},
            manager=manager,
            emit=emit,
            mode="ssh",
            intent=intent,
            task_title=task_title,
            task_issue=task_issue,
        )

    if vm_inventory:
        await _run_ssh_vm_inventory(
            results,
            task_id=task_id,
            label=label,
            server=SSH_MCP,
            manager=manager,
            emit=emit,
            mode="ssh",
            intent=intent,
            task_title=task_title,
            task_issue=task_issue,
            max_diagnostics=max_diagnostics,
        )
    return True


async def _run_ssh_vm_inventory(
    results: list[dict],
    *,
    task_id: int,
    label: str,
    server: str,
    manager,
    emit: EmitFn,
    mode: str,
    intent: str,
    task_title: str,
    task_issue: str,
    max_diagnostics: int,
) -> None:
    """After `qm list`, fetch `qm config` for each selected guest."""
    qemu_output = ""
    for r in results:
        if r.get("tool") != "ssh_run":
            continue
        out = r.get("output") or ""
        if re.search(r"^\s*\d+\s+\S+\s+(running|stopped)", out, re.MULTILINE):
            qemu_output = out
            break

    vms = _vms_for_detail(_parse_qemu_vms(qemu_output), task_title, task_issue)
    for vm in vms[:15]:
        vmid = vm.get("vmid")
        if vmid is None or len(results) >= max_diagnostics:
            break
        await _append_tool_result(
            results,
            task_id=task_id,
            label=label,
            server=server,
            tool="ssh_run",
            args={"command": f"qm config {int(vmid)} 2>/dev/null || true"},
            manager=manager,
            emit=emit,
            mode=mode,
            intent=intent,
            task_title=task_title,
            task_issue=task_issue,
        )


async def run_proxmox_diagnostics(
    task_id: int,
    device: dict,
    manager,
    emit: EmitFn,
    *,
    max_diagnostics: int = 6,
    task_title: str = "",
    task_issue: str = "",
) -> list[dict]:
    """Connect via Proxmox API when configured; otherwise or on failure use SSH."""
    results: list[dict] = []
    label = device.get("catalog_key") or device.get("name") or "proxmox"
    ctx = enrich_proxmox_device(device)
    vm_inventory = wants_vm_inventory(task_title, task_issue)
    intent = "vm inventory" if vm_inventory else "health"
    budget = MAX_VM_INVENTORY_CALLS if vm_inventory else max_diagnostics

    mode: str | None = None
    server: str | None = None
    diagnostics: list[tuple[str, Callable[[dict], dict]]] = []

    await emit(task_id, "connect", f"Connecting to {label}...", {"device": label})

    if proxmox_api_configured(ctx):
        await emit(
            task_id,
            "connect",
            f"{label}: trying Proxmox API ({ctx['proxmox_api_url']})...",
            {"device": label, "mode": "api"},
        )
        if await _try_proxmox_api(manager, ctx):
            mode, server, diagnostics = "api", PROXMOX_MCP, PROXMOX_API_DIAGNOSTICS
            await emit(
                task_id,
                "connect",
                f"{label}: Proxmox API connected ({intent})",
                {"device": label, "mode": "api", "intent": intent},
            )
        else:
            await emit(
                task_id,
                "connect",
                f"{label}: Proxmox API unavailable — falling back to SSH",
                {"device": label, "mode": "api_failed"},
            )

    if mode is None:
        await emit(
            task_id,
            "connect",
            f"{label}: trying SSH on port {ctx.get('port') or 22}...",
            {"device": label, "mode": "ssh"},
        )
        if await _try_proxmox_ssh(manager, ctx):
            mode, server, diagnostics = "ssh", SSH_MCP, PROXMOX_SSH_DIAGNOSTICS
            await emit(
                task_id,
                "connect",
                f"{label}: SSH connected ({intent})",
                {"device": label, "mode": "ssh", "intent": intent},
            )
        else:
            await emit(
                task_id,
                "connect",
                f"{label}: connect failed (no working API token or SSH credentials)",
                {"device": label},
            )
            return results

    for tool, builder in diagnostics[:budget]:
        args = builder(ctx)
        res, tool_ok = await invoke_tool(
            task_id,
            label,
            server,
            tool,
            args,
            manager,
            emit,
            task_title=task_title,
            task_issue=task_issue,
        )
        results.append(
            {
                "device": label,
                "tool": tool,
                "ok": tool_ok,
                "output": res.get("text", ""),
                "connect_mode": mode,
                "intent": intent,
            }
        )

    if mode == "api":
        node = _first_proxmox_node(next((r["output"] for r in results if r["tool"] == "proxmox_list_nodes"), ""))
        if node:
            await _run_api_node_tools(
                results,
                task_id=task_id,
                label=label,
                server=server,
                node=node,
                manager=manager,
                emit=emit,
                mode=mode,
                intent=intent,
                task_title=task_title,
                task_issue=task_issue,
                max_diagnostics=budget,
                vm_inventory=vm_inventory,
            )
            if _api_guest_inventory_empty(results):
                await _run_ssh_inventory_fallback(
                    results,
                    task_id=task_id,
                    label=label,
                    device=device,
                    manager=manager,
                    emit=emit,
                    intent=intent,
                    task_title=task_title,
                    task_issue=task_issue,
                    max_diagnostics=budget,
                    vm_inventory=vm_inventory,
                )
    elif mode == "ssh" and vm_inventory:
        await _run_ssh_vm_inventory(
            results,
            task_id=task_id,
            label=label,
            server=server,
            manager=manager,
            emit=emit,
            mode=mode,
            intent=intent,
            task_title=task_title,
            task_issue=task_issue,
            max_diagnostics=budget,
        )

    return results