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

"""Cull-distance ops — batch set/inspect Desired Max Draw Distance on ISM/HISM/StaticMesh components.

Encapsulates the gotchas learned the hard way:
  * CachedMaxDrawDistance is READ-ONLY  -> we only write LDMaxDrawDistance (+ SetCullDistance on live instances).
  * SubobjectDataHandle must be unwrapped via k2_find_subobject_data_from_handle (UE 5.7).
  * target='asset_template' edits the Blueprint asset (default for ALL instances, survives reconstruct);
    target='level_instance' edits placed actors only (fast, but lost on BP recompile / PLA repack).
  * only_zero keeps it idempotent and never clobbers hand-tuned non-zero distances.

Ops:
  audit  {scope?, asset_paths?, component_types?, target?}                 — read-only report, no mutation
  apply  {distance, scope?, asset_paths?, component_types?, only_zero?, target?, save?:true}
  verify {...same as audit}                                               — alias of audit (post-apply check)

  scope:           'level' (default) | 'assets' | 'selection'
  target:          'asset_template' (default) | 'level_instance'
  component_types: subset of ['ISM','HISM','StaticMesh','Primitive']  (default ['ISM'])
  distance:        cm; 0 disables culling
"""
from __future__ import annotations

import json
import os
import traceback

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


def _eas():  return unreal.get_editor_subsystem(unreal.EditorActorSubsystem)
def _sub():  return unreal.get_engine_subsystem(unreal.SubobjectDataSubsystem)


def _type_classes(args):
    types = args.get("component_types") or ["ISM"]
    mapping = {
        "ISM":        unreal.InstancedStaticMeshComponent,
        "HISM":       unreal.HierarchicalInstancedStaticMeshComponent,
        "StaticMesh": unreal.StaticMeshComponent,
        "Primitive":  unreal.PrimitiveComponent,
    }
    out = []
    for t in types:
        cls = mapping.get(t)
        if cls is None:
            raise ValueError(f"unknown component_type {t!r}; allowed: {list(mapping)}")
        out.append(cls)
    return out


def _matches(obj, classes):
    return any(isinstance(obj, c) for c in classes)


def _asset_path_from_actor(actor):
    gen = actor.get_class().get_path_name()
    asset = gen.split(".")[0]
    return asset[:-2] if asset.endswith("_C") else asset


# ---- target='asset_template' : edit Blueprint component templates -----------

def _template_components(bp, classes):
    """Returns deduped list of (name, component_template) for matching types."""
    handles = _sub().k2_gather_subobject_data_for_blueprint(bp)
    seen, out = set(), []
    for h in handles:
        data = _sub().k2_find_subobject_data_from_handle(h)
        obj = unreal.SubobjectDataBlueprintFunctionLibrary.get_object(data)
        if obj is None or not _matches(obj, classes):
            continue
        nm = obj.get_name()
        if nm in seen:
            continue
        seen.add(nm)
        out.append((nm, obj))
    return out


def _collect_assets(args):
    """Returns dict {asset_path: blueprint_asset} for the requested scope."""
    ael = unreal.EditorAssetLibrary
    paths = {}
    scope = args.get("scope") or "level"
    if scope == "assets":
        for p in (args.get("asset_paths") or []):
            paths[p.split(".")[0]] = None
    else:  # level / selection -> discover from placed actors
        classes = _type_classes(args)
        actors = (_eas().get_selected_level_actors() if scope == "selection"
                  else _eas().get_all_level_actors())
        for a in actors:
            if any(a.get_components_by_class(c) for c in classes):
                paths[_asset_path_from_actor(a)] = None
    result = {}
    for p in paths:
        result[p] = ael.load_asset(p)
    return result


def _run_assets(args, mutate):
    classes = _type_classes(args)
    distance = float(args.get("distance", 0.0))
    only_zero = args.get("only_zero", True)
    save = args.get("save", True)
    rows, total_set = [], 0
    for path, bp in _collect_assets(args).items():
        if bp is None:
            rows.append({"asset": path, "error": "load failed"}); continue
        comps = _template_components(bp, classes)
        vals, set_n = [], 0
        for _nm, c in comps:
            cur = float(c.get_editor_property("ld_max_draw_distance"))
            if mutate:
                if not (only_zero and abs(cur) > 0.001):
                    c.set_editor_property("ld_max_draw_distance", distance)
                    set_n += 1; cur = distance
            vals.append(round(cur, 1))
        if mutate and set_n:
            unreal.BlueprintEditorLibrary.compile_blueprint(bp)
            if save:
                unreal.EditorAssetLibrary.save_asset(path)
            total_set += 1
        rows.append({
            "asset": path, "components": len(comps),
            "zeros": sum(1 for v in vals if abs(v) < 0.001),
            "distinct": sorted(set(vals)), "set": set_n,
        })
    return {"ok": True, "target": "asset_template",
            "assets_modified": total_set, "rows": rows}


# ---- target='level_instance' : edit placed actors --------------------------

def _run_instances(args, mutate):
    classes = _type_classes(args)
    distance = float(args.get("distance", 0.0))
    only_zero = args.get("only_zero", True)
    scope = args.get("scope") or "level"
    actors = (_eas().get_selected_level_actors() if scope == "selection"
              else _eas().get_all_level_actors())
    rows, changed_actors, changed_comps = [], 0, 0
    for a in actors:
        comps, seen = [], set()
        for c in classes:
            for comp in a.get_components_by_class(c):
                if comp.get_name() not in seen:
                    seen.add(comp.get_name()); comps.append(comp)
        if not comps:
            continue
        vals, set_n = [], 0
        for comp in comps:
            cur = float(comp.get_editor_property("ld_max_draw_distance"))
            if mutate and not (only_zero and abs(cur) > 0.001):
                comp.modify()
                comp.set_editor_property("ld_max_draw_distance", distance)
                comp.set_cull_distance(distance)  # updates cached + re-registers
                set_n += 1; cur = distance
            vals.append(round(cur, 1))
        if mutate and set_n:
            a.modify(); changed_actors += 1; changed_comps += set_n
        rows.append({"actor": a.get_actor_label(), "components": len(comps),
                     "zeros": sum(1 for v in vals if abs(v) < 0.001),
                     "distinct": sorted(set(vals)), "set": set_n})
    return {"ok": True, "target": "level_instance",
            "actors_modified": changed_actors, "components_set": changed_comps,
            "rows": rows, "note": "instance edits are lost on BP recompile / PLA repack"}


# ---- op table --------------------------------------------------------------

def _dispatch(args, mutate):
    target = args.get("target") or "asset_template"
    if target == "asset_template":
        return _run_assets(args, mutate)
    if target == "level_instance":
        return _run_instances(args, mutate)
    return {"ok": False, "error": f"unknown target {target!r}; use 'asset_template' or 'level_instance'"}


def audit(args):
    return _dispatch(args, mutate=False)


def verify(args):
    return _dispatch(args, mutate=False)


def apply(args):
    if "distance" not in args:
        return {"ok": False, "error": "apply requires 'distance' (cm)"}
    return _dispatch(args, mutate=True)


OPS = {"audit": audit, "verify": verify, "apply": apply}


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-CULL] " + txt[:1500])
    except Exception:
        pass
    return txt