unreal-engine-mcp-system-plugin/ue_side/selection.py

"""Selection-aware Blueprint graph ops.

Ops:
  get_context      {bp_path, function?}
  apply_relative   {bp_path, function?, placement?, gap_x?, gap_y?,
                    delete_selection?, nodes:[], edges:[]}

This module keeps partial graph generation out of the full-graph replace path:
LLMs can inspect the user's selected nodes, then append or replace only that
local area.
"""
from __future__ import annotations

import json
import os
import traceback

import graph_layout

try:
    import unreal  # type: ignore
except ImportError:
    unreal = None


def _lib():
    L = getattr(unreal, "MCPGraphLibrary", None)
    if L is None:
        raise RuntimeError("MCPGraphLibrary missing - rebuild C++ plugin")
    return L


def _load_bp(path: str):
    a = unreal.EditorAssetLibrary.load_asset(path)
    if not isinstance(a, unreal.Blueprint):
        raise RuntimeError(f"not a blueprint: {path!r}")
    return a


def _graph(bp, fn: str | None):
    L = _lib()
    if fn in (None, "", "EventGraph"):
        return L.find_event_graph(bp)
    return L.find_function_graph(bp, fn)


def _selected(bp) -> dict:
    raw = _lib().get_selected_nodes_json(bp)
    data = json.loads(raw)
    nodes = data.get("selected") or []
    bbox = _bbox(nodes)
    focused_graph = data.get("focused_graph")
    draft_zones = _draft_zones(bp, focused_graph)
    return {
        "focused_graph": focused_graph,
        "selected": nodes,
        "selected_count": len(nodes),
        "bbox": bbox,
        "anchors": _anchors(bbox),
        "draft_zones": draft_zones,
    }


def _draft_zones(bp, function_name: str | None) -> list[dict]:
    graph = _graph(bp, function_name or "EventGraph")
    if graph is None:
        return []
    raw = _lib().dump_graph_to_json(graph)
    data = json.loads(raw)
    zones = []
    for node in data.get("nodes") or []:
        if node.get("class") != "EdGraphNode_Comment":
            continue
        comment = str(node.get("comment") or node.get("title") or "")
        if comment.strip().lower() not in ("mcp zone", "mcp draft zone"):
            continue
        zone = dict(node)
        zone["bbox"] = {
            "min_x": float(node.get("x", 0)),
            "min_y": float(node.get("y", 0)),
            "max_x": float(node.get("x", 0)) + float(node.get("w", 0)),
            "max_y": float(node.get("y", 0)) + float(node.get("h", 0)),
            "width": float(node.get("w", 0)),
            "height": float(node.get("h", 0)),
            "center_x": float(node.get("x", 0)) + (float(node.get("w", 0)) / 2.0),
            "center_y": float(node.get("y", 0)) + (float(node.get("h", 0)) / 2.0),
        }
        zones.append(zone)
    return zones


def _bbox(nodes: list[dict]) -> dict | None:
    if not nodes:
        return None
    min_x = min(float(n.get("x", 0)) for n in nodes)
    min_y = min(float(n.get("y", 0)) for n in nodes)
    max_x = max(float(n.get("x", 0)) + float(n.get("w", 220)) for n in nodes)
    max_y = max(float(n.get("y", 0)) + float(n.get("h", 120)) for n in nodes)
    return {
        "min_x": min_x,
        "min_y": min_y,
        "max_x": max_x,
        "max_y": max_y,
        "width": max_x - min_x,
        "height": max_y - min_y,
        "center_x": min_x + ((max_x - min_x) / 2.0),
        "center_y": min_y + ((max_y - min_y) / 2.0),
    }


def _anchors(bbox: dict | None) -> dict:
    if not bbox:
        return {
            "right": [0, 0],
            "below": [0, 320],
            "center": [0, 0],
        }
    return {
        "right": [bbox["max_x"] + 320, bbox["min_y"]],
        "below": [bbox["min_x"], bbox["max_y"] + 220],
        "center": [bbox["center_x"], bbox["center_y"]],
    }


def get_context(args: dict) -> dict:
    bp = _load_bp(args["bp_path"])
    ctx = _selected(bp)
    return {"ok": True, "bp_path": args["bp_path"], **ctx}


def apply_relative(args: dict) -> dict:
    import importlib
    import apply_graph

    importlib.reload(apply_graph)

    bp = _load_bp(args["bp_path"])
    ctx = _selected(bp)
    function_name = args.get("function") or ctx.get("focused_graph") or "EventGraph"
    target_graph = _graph(bp, function_name)
    if target_graph is None:
        return {"ok": False, "error": f"graph not found: {function_name!r}", "selection": ctx}

    placement = args.get("placement") or "right"
    gap_x = float(args.get("gap_x", 0))
    gap_y = float(args.get("gap_y", 0))

    zone_guid = args.get("zone_guid") or ""
    zone = None
    for candidate in ctx.get("draft_zones") or []:
        if zone_guid and candidate.get("guid") == zone_guid:
            zone = candidate
            break
    if zone is None and not ctx.get("selected") and ctx.get("draft_zones"):
        zone = ctx["draft_zones"][-1]

    if zone is not None:
        bbox = zone["bbox"]
        zone_anchors = {
            "right": [bbox["min_x"] + 80, bbox["min_y"] + 80],
            "below": [bbox["min_x"] + 80, bbox["min_y"] + 80],
            "center": [bbox["center_x"], bbox["center_y"]],
        }
        anchor = zone_anchors.get(placement) or zone_anchors["right"]
    else:
        anchor = ctx["anchors"].get(placement)
        if anchor is None:
            anchor = ctx["anchors"]["right"]
    ax, ay = float(anchor[0]) + gap_x, float(anchor[1]) + gap_y

    relative_nodes = []
    for spec in args.get("nodes") or []:
        n = dict(spec)
        x, y = n.get("position") or [0, 0]
        n["position"] = [float(x) + ax, float(y) + ay]
        relative_nodes.append(n)

    nodes = graph_layout.normalize_specs(
        relative_nodes,
        args.get("edges") or [],
        domain="blueprint",
    ) if args.get("auto_layout", True) else relative_nodes

    deleted = []
    if args.get("delete_selection"):
        for n in ctx.get("selected") or []:
            guid = n.get("guid")
            if guid and _lib().delete_node(target_graph, guid):
                deleted.append(guid)

    graph = {
        "bp_path": args["bp_path"],
        "function": function_name,
        "mode": "append",
        "nodes": nodes,
        "edges": args.get("edges") or [],
    }
    result = apply_graph.apply(graph)
    return {
        "ok": bool(result.get("ok")),
        "bp_path": args["bp_path"],
        "function": function_name,
        "placement": placement,
        "anchor": [ax, ay],
        "zone_guid": zone.get("guid") if zone else None,
        "deleted_selection": deleted,
        "selection": ctx,
        "applied": result,
    }


OPS = {
    "get_context": get_context,
    "apply_relative": apply_relative,
}


def entrypoint(payload_json: str) -> str:
    try:
        p = json.loads(payload_json)
    except Exception as e:  # noqa: BLE001
        return _tee({"ok": False, "error": f"bad json: {e}"}, {})
    fn = OPS.get(p.get("op"))
    if not fn:
        return _tee({"ok": False, "error": f"unknown op {p.get('op')!r}", "available": list(OPS)}, p)
    try:
        return _tee(fn(p), p)
    except Exception as e:  # noqa: BLE001
        return _tee({"ok": False, "error": str(e), "traceback": traceback.format_exc()}, p)


def _tee(result: dict, p: dict) -> str:
    txt = json.dumps(result)
    try:
        if unreal is not None:
            sd = unreal.Paths.project_saved_dir() + "MCP/"
            os.makedirs(sd, exist_ok=True)
            rid = p.get("__rid") or "last"
            with open(sd + f"{rid}.json", "w", encoding="utf-8") as f:
                f.write(txt)
            with open(sd + "last.json", "w", encoding="utf-8") as f:
                f.write(txt)
            unreal.log("[MCP-SELECTION] " + txt[:1500])
    except Exception:
        pass
    return txt