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

"""Enhanced Input ops — Input Actions, Input Mapping Contexts.

Ops:
  create_input_action          {path, value_type?:'Bool|Axis1D|Axis2D|Axis3D'}
  create_input_mapping_context {path}
  add_mapping                  {imc_path, action_path, key, modifiers?:[class_path], triggers?:[class_path]}
  remove_mapping               {imc_path, index}
  list_mappings                {imc_path}
"""
from __future__ import annotations

import json
import os
import traceback

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


def _at(): return unreal.AssetToolsHelpers.get_asset_tools()
def _load(p): return unreal.EditorAssetLibrary.load_asset(p)
def _split(p):
    folder, name = p.rsplit("/", 1)
    return folder, name


VALUE_TYPES = {
    "Bool":   "Boolean",
    "Axis1D": "Axis1D",
    "Axis2D": "Axis2D",
    "Axis3D": "Axis3D",
}


def create_input_action(args):
    folder, name = _split(args["path"])
    factory_cls = getattr(unreal, "InputActionFactory", None)
    asset_cls = getattr(unreal, "InputAction", None)
    if factory_cls is None or asset_cls is None:
        return {"ok": False, "error": "Enhanced Input plugin not enabled"}
    asset = _at().create_asset(name, folder, asset_cls, factory_cls())
    if asset is None:
        return {"ok": False, "error": "create returned None"}
    vt = args.get("value_type") or "Bool"
    enum_name = VALUE_TYPES.get(vt)
    if enum_name:
        try:
            enum_obj = getattr(unreal.InputActionValueType, enum_name)
            asset.set_editor_property("value_type", enum_obj)
        except Exception:
            pass
    unreal.EditorAssetLibrary.save_asset(asset.get_path_name())
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def create_input_mapping_context(args):
    folder, name = _split(args["path"])
    factory_cls = getattr(unreal, "InputMappingContextFactory", None)
    asset_cls = getattr(unreal, "InputMappingContext", None)
    if factory_cls is None or asset_cls is None:
        return {"ok": False, "error": "Enhanced Input plugin not enabled"}
    asset = _at().create_asset(name, folder, asset_cls, factory_cls())
    if asset is None:
        return {"ok": False, "error": "create returned None"}
    unreal.EditorAssetLibrary.save_asset(asset.get_path_name())
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def add_mapping(args):
    imc = _load(args["imc_path"])
    ia = _load(args["action_path"])
    if imc is None or ia is None:
        return {"ok": False, "error": "imc or action not found"}
    key = unreal.InputCoreLibrary.input_key_from_string(args["key"]) \
        if hasattr(unreal, "InputCoreLibrary") and hasattr(unreal.InputCoreLibrary, "input_key_from_string") \
        else unreal.Key(args["key"])
    mapping_cls = getattr(unreal, "EnhancedActionKeyMapping", None)
    if mapping_cls is None:
        try:
            imc.map_key(ia, key)
            unreal.EditorAssetLibrary.save_asset(imc.get_path_name())
            return {"ok": True, "note": "added via map_key()"}
        except Exception as e:
            return {"ok": False, "error": f"map_key failed: {e}"}
    m = mapping_cls()
    m.set_editor_property("action", ia)
    m.set_editor_property("key", key)
    # Modifiers/triggers (class CDOs):
    for cp in (args.get("modifiers") or []):
        mc = unreal.load_class(None, cp)
        if mc is not None:
            try:
                inst = unreal.new_object(mc)
                mods = list(m.get_editor_property("modifiers") or [])
                mods.append(inst)
                m.set_editor_property("modifiers", mods)
            except Exception:
                pass
    for cp in (args.get("triggers") or []):
        tc = unreal.load_class(None, cp)
        if tc is not None:
            try:
                inst = unreal.new_object(tc)
                trigs = list(m.get_editor_property("triggers") or [])
                trigs.append(inst)
                m.set_editor_property("triggers", trigs)
            except Exception:
                pass
    mappings = list(imc.get_editor_property("mappings") or [])
    mappings.append(m)
    imc.set_editor_property("mappings", mappings)
    unreal.EditorAssetLibrary.save_asset(imc.get_path_name())
    return {"ok": True, "index": len(mappings) - 1}


def remove_mapping(args):
    imc = _load(args["imc_path"])
    if imc is None:
        return {"ok": False, "error": "imc not found"}
    mappings = list(imc.get_editor_property("mappings") or [])
    idx = int(args["index"])
    if idx < 0 or idx >= len(mappings):
        return {"ok": False, "error": "index out of range"}
    mappings.pop(idx)
    imc.set_editor_property("mappings", mappings)
    unreal.EditorAssetLibrary.save_asset(imc.get_path_name())
    return {"ok": True}


def list_mappings(args):
    imc = _load(args["imc_path"])
    if imc is None:
        return {"ok": False, "error": "imc not found"}
    out = []
    for i, m in enumerate(imc.get_editor_property("mappings") or []):
        try:
            a = m.get_editor_property("action")
            k = m.get_editor_property("key")
            out.append({
                "index": i,
                "action": a.get_path_name().split(".")[0] if a else None,
                "key": str(k),
            })
        except Exception:
            pass
    return {"ok": True, "mappings": out}


OPS = {
    "create_input_action":          create_input_action,
    "create_input_mapping_context": create_input_mapping_context,
    "add_mapping":                  add_mapping,
    "remove_mapping":               remove_mapping,
    "list_mappings":                list_mappings,
}


def entrypoint(payload_json):
    try:
        p = json.loads(payload_json)
    except Exception as e:
        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:
        return _tee({"ok": False, "error": str(e), "traceback": traceback.format_exc()}, p)


def _tee(result, p):
    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-INPUT] " + txt[:1500])
    except Exception:
        pass
    return txt