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

"""Niagara authoring ops. Bridges to unreal.MCPNiagaraLibrary (C++) where
editor-only APIs are required, falls back to pure Python where possible.

Ops:
  create_system        {path}
  create_emitter_asset {path, template_emitter?}             — create empty UNiagaraEmitter asset
  add_emitter          {system, emitter_asset, name?}        — add emitter to system, returns handle name
  remove_emitter       {system, emitter_name}
  list_emitters        {system}
  list_modules         {system, emitter_name?, group}        — group: EmitterSpawn|EmitterUpdate|ParticleSpawn|ParticleUpdate|SystemSpawn|SystemUpdate
  add_module           {system, emitter_name?, group, module_script}
  remove_module        {system, emitter_name?, group, module_name}
  add_renderer         {system, emitter_name, renderer_class}     — e.g. NiagaraSpriteRendererProperties
  remove_renderer      {system, emitter_name, renderer_class}
  compile              {system}
  set_module_input     STUB — see NOTES at bottom.
"""
from __future__ import annotations

import json
import os
import traceback

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


# ---------- loaders ----------

def _load_asset(path: str):
    if not path:
        return None
    return unreal.EditorAssetLibrary.load_asset(path)


def _load_class(p: str):
    if not p:
        return None
    if "/" in p:
        return unreal.load_class(None, p)
    # Try common Niagara paths
    for prefix in ("/Script/Niagara.", "/Script/NiagaraEditor."):
        c = unreal.load_class(None, prefix + p)
        if c is not None:
            return c
    return unreal.load_class(None, p)


def _lib():
    if not hasattr(unreal, "MCPNiagaraLibrary"):
        raise RuntimeError("MCPNiagaraLibrary not found — rebuild the UEBlueprintMCPEditor C++ module")
    return unreal.MCPNiagaraLibrary


# ---------- ops ----------

def create_system(p):
    path = p["path"]
    ok = _lib().create_niagara_system(path)
    return {"ok": bool(ok), "path": path}


def create_emitter_asset(p):
    """Create an empty UNiagaraEmitter asset. Niagara typically doesn't support
    fully empty emitters via UI — users start from a template. We create the bare
    asset; for fully usable emitters duplicate an existing template via asset_op.
    """
    path = p["path"]
    package_name, _, asset_name = path.rpartition("/")
    factory = unreal.NiagaraEmitterFactoryNew() if hasattr(unreal, "NiagaraEmitterFactoryNew") else None
    tools = unreal.AssetToolsHelpers.get_asset_tools()
    if factory is None:
        return {"ok": False, "error": "NiagaraEmitterFactoryNew not exposed to Python — duplicate a template emitter via asset_op instead"}
    asset = tools.create_asset(asset_name, package_name, unreal.NiagaraEmitter, factory)
    return {"ok": asset is not None, "path": path}


def add_emitter(p):
    system = _load_asset(p["system"])
    emitter = _load_asset(p["emitter_asset"])
    if system is None or emitter is None:
        return {"ok": False, "error": "system or emitter_asset not found"}
    name = p.get("name", "")
    handle_name = _lib().add_emitter_to_system(system, emitter, name)
    if not handle_name:
        return {"ok": False, "error": "AddEmitterToSystem returned empty"}
    unreal.EditorAssetLibrary.save_asset(p["system"])
    return {"ok": True, "handle_name": handle_name}


def remove_emitter(p):
    system = _load_asset(p["system"])
    if system is None:
        return {"ok": False, "error": "system not found"}
    ok = _lib().remove_emitter_from_system(system, p["emitter_name"])
    if ok:
        unreal.EditorAssetLibrary.save_asset(p["system"])
    return {"ok": bool(ok)}


def list_emitters(p):
    system = _load_asset(p["system"])
    if system is None:
        return {"ok": False, "error": "system not found"}
    names = list(_lib().get_emitter_handle_names(system))
    return {"ok": True, "emitters": names}


def list_modules(p):
    system = _load_asset(p["system"])
    if system is None:
        return {"ok": False, "error": "system not found"}
    modules = list(_lib().list_modules(system, p.get("emitter_name", ""), p["group"]))
    return {"ok": True, "modules": modules}


def add_module(p):
    system = _load_asset(p["system"])
    script = _load_asset(p["module_script"])
    if system is None or script is None:
        return {"ok": False, "error": "system or module_script not found"}
    name = _lib().add_module_to_stack(
        system,
        p.get("emitter_name", ""),
        p["group"],
        script,
        int(p.get("index", -1)),
    )
    if not name:
        return {"ok": False, "error": "AddModuleToStack returned empty (wrong group or script not a module?)"}
    unreal.EditorAssetLibrary.save_asset(p["system"])
    return {"ok": True, "module_name": name}


def remove_module(p):
    system = _load_asset(p["system"])
    if system is None:
        return {"ok": False, "error": "system not found"}
    ok = _lib().remove_module_from_stack(
        system,
        p.get("emitter_name", ""),
        p["group"],
        p["module_name"],
    )
    if ok:
        unreal.EditorAssetLibrary.save_asset(p["system"])
    return {"ok": bool(ok)}


def add_renderer(p):
    system = _load_asset(p["system"])
    if system is None:
        return {"ok": False, "error": "system not found"}
    cls = _load_class(p["renderer_class"])
    if cls is None:
        return {"ok": False, "error": f"renderer class {p['renderer_class']!r} not found"}
    result = _lib().add_renderer_to_emitter(system, p["emitter_name"], cls)
    if not result:
        return {"ok": False, "error": "AddRendererToEmitter returned empty"}
    unreal.EditorAssetLibrary.save_asset(p["system"])
    return {"ok": True, "renderer": result}


def remove_renderer(p):
    system = _load_asset(p["system"])
    if system is None:
        return {"ok": False, "error": "system not found"}
    cls = _load_class(p["renderer_class"])
    if cls is None:
        return {"ok": False, "error": f"renderer class {p['renderer_class']!r} not found"}
    n = int(_lib().remove_renderers_from_emitter(system, p["emitter_name"], cls))
    if n:
        unreal.EditorAssetLibrary.save_asset(p["system"])
    return {"ok": True, "removed": n}


def assign_to_actor(p):
    """Assign a UNiagaraSystem to the NiagaraComponent of a NiagaraActor in the
    current level. Identifies the actor by its label (use `name` field).
    """
    system_path = p["system"]
    actor_name = p["name"]
    system = _load_asset(system_path)
    if system is None:
        return {"ok": False, "error": f"system {system_path!r} not found"}
    ess = unreal.get_editor_subsystem(unreal.EditorActorSubsystem)
    target = None
    for a in ess.get_all_level_actors():
        if a and (a.get_actor_label() == actor_name or a.get_name() == actor_name):
            target = a
            break
    if target is None:
        return {"ok": False, "error": f"actor {actor_name!r} not found"}
    comp = target.get_component_by_class(unreal.NiagaraComponent)
    if comp is None:
        return {"ok": False, "error": "actor has no NiagaraComponent"}
    comp.set_asset(system)
    comp.activate(reset=True)
    return {"ok": True, "actor": actor_name, "system": system_path}


def compile_system(p):
    system = _load_asset(p["system"])
    if system is None:
        return {"ok": False, "error": "system not found"}
    ok = _lib().request_compile(system)
    return {"ok": bool(ok)}


def set_module_input(p):
    return {
        "ok": False,
        "error": (
            "set_module_input not implemented yet — requires "
            "FNiagaraStackFunctionInputBinder + NiagaraClipboard plumbing. "
            "Workaround: open the system in editor and use UI, or extend "
            "MCPNiagaraLibrary with SetModuleInput(scalar|vector|enum) wrappers."
        ),
    }


OPS = {
    "create_system":        create_system,
    "create_emitter_asset": create_emitter_asset,
    "add_emitter":          add_emitter,
    "remove_emitter":       remove_emitter,
    "list_emitters":        list_emitters,
    "list_modules":         list_modules,
    "add_module":           add_module,
    "remove_module":        remove_module,
    "add_renderer":         add_renderer,
    "remove_renderer":      remove_renderer,
    "compile":              compile_system,
    "assign_to_actor":      assign_to_actor,
    "set_module_input":     set_module_input,
}


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}"}, {})
    op = p.get("op")
    fn = OPS.get(op)
    if not fn:
        return _tee({"ok": False, "error": f"unknown op {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 + "last.json", "w", encoding="utf-8") as f:
                f.write(txt)
            if rid != "last":
                with open(sd + f"{rid}.json", "w", encoding="utf-8") as f:
                    f.write(txt)
            unreal.log("[MCP-NIAGARA] " + txt[:1500])
    except Exception:
        pass
    return txt