Unreal-mcp/src/unreal_mcp/tools/scene.py

"""Scene Manipulation tools for Unreal MCP Server."""

from typing import Any

from unreal_mcp.core.unreal_connection import UnrealConnection
from unreal_mcp.utils.logger import get_logger
from unreal_mcp.utils.validation import validate_class_name, validate_location, validate_rotation

logger = get_logger(__name__)


async def spawn_actor(
    conn: UnrealConnection,
    class_name: str,
    location: list[float],
    rotation: list[float] | None = None,
    properties: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Spawn an actor in the current level.

    Args:
        conn: Unreal connection instance
        class_name: Name of the class to spawn
        location: Spawn location [X, Y, Z]
        rotation: Optional spawn rotation [Pitch, Yaw, Roll]
        properties: Optional initial property values

    Returns:
        Dictionary with spawn result
    """
    # Validate inputs
    try:
        class_name = validate_class_name(class_name)
        loc = validate_location(location)
        rot = validate_rotation(rotation) if rotation else (0, 0, 0)
    except ValueError as e:
        return {"success": False, "error": str(e), "code": "VALIDATION_ERROR"}

    # Check connection
    if not conn.is_connected and not conn.connect():
        return {
            "success": False,
            "error": "Cannot spawn actor: Editor not connected",
            "code": "NOT_CONNECTED",
        }

    # Build spawn script
    props_code = ""
    if properties:
        for prop_name, prop_value in properties.items():
            if isinstance(prop_value, str):
                props_code += f"    actor.set_editor_property('{prop_name}', '{prop_value}')\n"
            else:
                props_code += f"    actor.set_editor_property('{prop_name}', {prop_value})\n"

    script = f"""
import unreal

# Find the class
actor_class = unreal.EditorAssetLibrary.find_asset_data('/Game/**/{class_name}*').get_asset()
if not actor_class:
    # Try as a native class
    actor_class = getattr(unreal, '{class_name}', None)

if actor_class:
    location = unreal.Vector({loc[0]}, {loc[1]}, {loc[2]})
    rotation = unreal.Rotator({rot[0]}, {rot[1]}, {rot[2]})

    actor = unreal.EditorLevelLibrary.spawn_actor_from_class(
        actor_class,
        location,
        rotation
    )

    if actor:
{props_code if props_code else "        pass"}
        result = {{
            'success': True,
            'actor_id': actor.get_name(),
            'location': [{loc[0]}, {loc[1]}, {loc[2]}]
        }}
    else:
        result = {{'success': False, 'error': 'Failed to spawn actor'}}
else:
    result = {{'success': False, 'error': 'Class not found: {class_name}'}}

print(result)
"""

    response = await conn.execute(script)

    if response.get("success"):
        data = response.get("data")
        if isinstance(data, dict):
            return data
        return {"success": True, "data": data}

    return response


async def get_scene_hierarchy(
    conn: UnrealConnection,
    include_components: bool = False,
) -> dict[str, Any]:
    """Get the hierarchy of actors in the current level.

    Args:
        conn: Unreal connection instance
        include_components: Include component details

    Returns:
        Dictionary with level name and actors list
    """
    if not conn.is_connected and not conn.connect():
        return {
            "success": False,
            "error": "Cannot get scene: Editor not connected",
            "code": "NOT_CONNECTED",
        }

    components_code = ""
    if include_components:
        components_code = """
            components = actor.get_components_by_class(unreal.ActorComponent)
            actor_info['components'] = [comp.get_name() for comp in components]
"""

    script = f"""
import unreal

actors = unreal.EditorLevelLibrary.get_all_level_actors()
level = unreal.EditorLevelLibrary.get_editor_world()

result = {{
    'level': level.get_name() if level else 'Unknown',
    'actors': []
}}

for actor in actors:
    location = actor.get_actor_location()
    actor_info = {{
        'id': actor.get_name(),
        'class': actor.get_class().get_name(),
        'location': [location.x, location.y, location.z],
    }}
{components_code}
    result['actors'].append(actor_info)

print(result)
"""

    response = await conn.execute(script)

    if response.get("success"):
        data = response.get("data")
        if isinstance(data, dict):
            return {"success": True, "data": data}
        return {"success": True, "data": data}

    return response


async def modify_actor_transform(
    conn: UnrealConnection,
    actor_id: str,
    location: list[float] | None = None,
    rotation: list[float] | None = None,
    scale: list[float] | None = None,
) -> dict[str, Any]:
    """Modify the transform of an actor.

    Args:
        conn: Unreal connection instance
        actor_id: ID or name of the actor to modify
        location: New location [X, Y, Z]
        rotation: New rotation [Pitch, Yaw, Roll]
        scale: New scale [X, Y, Z]

    Returns:
        Dictionary with modification result
    """
    if not conn.is_connected and not conn.connect():
        return {
            "success": False,
            "error": "Cannot modify actor: Editor not connected",
            "code": "NOT_CONNECTED",
        }

    # Build modification code
    mods = []
    if location:
        loc = validate_location(location)
        mods.append(f"actor.set_actor_location(unreal.Vector({loc[0]}, {loc[1]}, {loc[2]}), False, False)")
    if rotation:
        rot = validate_rotation(rotation)
        mods.append(f"actor.set_actor_rotation(unreal.Rotator({rot[0]}, {rot[1]}, {rot[2]}), False)")
    if scale:
        if len(scale) != 3:
            return {"success": False, "error": "Scale must have 3 components", "code": "VALIDATION_ERROR"}
        mods.append(f"actor.set_actor_scale3d(unreal.Vector({scale[0]}, {scale[1]}, {scale[2]}))")

    if not mods:
        return {"success": False, "error": "No transform changes specified", "code": "NO_CHANGES"}

    mods_code = "\n        ".join(mods)

    script = f"""
import unreal

actors = unreal.EditorLevelLibrary.get_all_level_actors()
actor = None

for a in actors:
    if a.get_name() == '{actor_id}':
        actor = a
        break

if actor:
    try:
        {mods_code}
        result = {{'success': True, 'actor_id': '{actor_id}'}}
    except Exception as e:
        result = {{'success': False, 'error': str(e)}}
else:
    result = {{'success': False, 'error': 'Actor not found: {actor_id}'}}

print(result)
"""

    response = await conn.execute(script)

    if response.get("success"):
        data = response.get("data")
        if isinstance(data, dict):
            return data
        return {"success": True, "data": data}

    return response