unreal-engine-mcp-system-plugin/Source/UEBlueprintMCPEditor/Private/MCPNiagaraLibrary.cpp

#include "MCPNiagaraLibrary.h"

#include "AssetToolsModule.h"
#include "IAssetTools.h"
#include "AssetRegistry/AssetData.h"
#include "Modules/ModuleManager.h"
#include "UObject/Package.h"
#include "UObject/UObjectGlobals.h"

#include "NiagaraSystem.h"
#include "NiagaraSystemFactoryNew.h"
#include "NiagaraEmitter.h"
#include "NiagaraEmitterHandle.h"
#include "NiagaraScript.h"
#include "NiagaraRendererProperties.h"
#include "NiagaraScriptSource.h"
#include "NiagaraGraph.h"
#include "NiagaraNodeOutput.h"
#include "NiagaraNodeFunctionCall.h"
#include "NiagaraCommon.h"
#include "ViewModels/Stack/NiagaraStackGraphUtilities.h"
#include "NiagaraEditorUtilities.h"

namespace
{
struct FStackLocation
{
    UNiagaraNodeOutput* OutputNode = nullptr;
    bool bValid() const { return OutputNode != nullptr; }
};

ENiagaraScriptUsage ResolveUsage(const FString& GroupName, bool& bOutIsSystem)
{
    bOutIsSystem = false;
    const FString G = GroupName.TrimStartAndEnd();
    if (G.Equals(TEXT("EmitterSpawn"),   ESearchCase::IgnoreCase)) return ENiagaraScriptUsage::EmitterSpawnScript;
    if (G.Equals(TEXT("EmitterUpdate"),  ESearchCase::IgnoreCase)) return ENiagaraScriptUsage::EmitterUpdateScript;
    if (G.Equals(TEXT("ParticleSpawn"),  ESearchCase::IgnoreCase)) return ENiagaraScriptUsage::ParticleSpawnScript;
    if (G.Equals(TEXT("ParticleUpdate"), ESearchCase::IgnoreCase)) return ENiagaraScriptUsage::ParticleUpdateScript;
    if (G.Equals(TEXT("SystemSpawn"),    ESearchCase::IgnoreCase)) { bOutIsSystem = true; return ENiagaraScriptUsage::SystemSpawnScript;  }
    if (G.Equals(TEXT("SystemUpdate"),   ESearchCase::IgnoreCase)) { bOutIsSystem = true; return ENiagaraScriptUsage::SystemUpdateScript; }
    return ENiagaraScriptUsage::ParticleSpawnScript;
}

UNiagaraNodeOutput* FindOutputNodeFromScript(UNiagaraScript* Script, ENiagaraScriptUsage Usage)
{
    if (!Script) return nullptr;
    UNiagaraScriptSourceBase* SourceBase = Script->GetLatestSource();
    UNiagaraScriptSource* Source = Cast<UNiagaraScriptSource>(SourceBase);
    if (!Source || !Source->NodeGraph) return nullptr;
    // FindOutputNode is not NIAGARAEDITOR_API; FindEquivalentOutputNode is.
    return Source->NodeGraph->FindEquivalentOutputNode(Usage, FGuid());
}

void GatherFunctionCallNodes(UNiagaraNodeOutput* OutputNode, TArray<UNiagaraNodeFunctionCall*>& Out)
{
    if (!OutputNode) return;
    UEdGraph* Graph = OutputNode->GetGraph();
    if (!Graph) return;
    for (UEdGraphNode* Node : Graph->Nodes)
    {
        if (UNiagaraNodeFunctionCall* FC = Cast<UNiagaraNodeFunctionCall>(Node))
        {
            // Filter to modules attached to this output's usage. GetCalledUsage()
            // returns the usage of the inner module script (usually Module).
            // Without param-map walking we list all function-calls in the graph
            // tied to the same output usage — good enough for list/remove.
            Out.Add(FC);
        }
    }
}

FNiagaraEmitterHandle* FindEmitterHandleMutable(UNiagaraSystem* System, const FString& EmitterName)
{
    if (!System) return nullptr;
    TArray<FNiagaraEmitterHandle>& Handles = System->GetEmitterHandles();
    for (FNiagaraEmitterHandle& H : Handles)
    {
        if (H.GetName().ToString().Equals(EmitterName, ESearchCase::IgnoreCase))
        {
            return &H;
        }
    }
    return nullptr;
}

FStackLocation LocateStack(UNiagaraSystem* System, const FString& EmitterName, const FString& GroupName)
{
    FStackLocation Loc;
    bool bIsSystem = false;
    const ENiagaraScriptUsage Usage = ResolveUsage(GroupName, bIsSystem);

    if (!System) return Loc;

    if (bIsSystem)
    {
        UNiagaraScript* Script = (Usage == ENiagaraScriptUsage::SystemSpawnScript)
            ? System->GetSystemSpawnScript()
            : System->GetSystemUpdateScript();
        Loc.OutputNode = FindOutputNodeFromScript(Script, Usage);
        return Loc;
    }

    FNiagaraEmitterHandle* Handle = FindEmitterHandleMutable(System, EmitterName);
    if (!Handle) return Loc;
    FVersionedNiagaraEmitter Versioned = Handle->GetInstance();
    FVersionedNiagaraEmitterData* Data = Versioned.GetEmitterData();
    if (!Data) return Loc;
    UNiagaraScript* Script = Data->GetScript(Usage, FGuid());
    Loc.OutputNode = FindOutputNodeFromScript(Script, Usage);
    return Loc;
}
}

bool UMCPNiagaraLibrary::CreateNiagaraSystem(const FString& PackagePath)
{
    if (PackagePath.IsEmpty()) return false;

    FString PackageName, AssetName;
    if (!PackagePath.Split(TEXT("/"), &PackageName, &AssetName, ESearchCase::IgnoreCase, ESearchDir::FromEnd))
    {
        return false;
    }

    IAssetTools& AssetTools = FModuleManager::LoadModuleChecked<FAssetToolsModule>("AssetTools").Get();
    UNiagaraSystemFactoryNew* Factory = NewObject<UNiagaraSystemFactoryNew>();
    UObject* Created = AssetTools.CreateAsset(AssetName, PackageName, UNiagaraSystem::StaticClass(), Factory);
    return Created != nullptr;
}

TArray<FString> UMCPNiagaraLibrary::GetEmitterHandleNames(UNiagaraSystem* System)
{
    TArray<FString> Out;
    if (!System) return Out;
    for (const FNiagaraEmitterHandle& H : System->GetEmitterHandles())
    {
        Out.Add(H.GetName().ToString());
    }
    return Out;
}

FString UMCPNiagaraLibrary::AddEmitterToSystem(UNiagaraSystem* System, UNiagaraEmitter* SourceEmitter, const FString& EmitterName)
{
    if (!System || !SourceEmitter) return FString();

    System->Modify();
    // Use the editor-utility path that performs full setup (merge, renderer
    // initialization, parameter refresh). The raw UNiagaraSystem::AddEmitterHandle
    // skips most of this and leaves the emitter non-functional in preview.
    const FGuid HandleId = FNiagaraEditorUtilities::AddEmitterToSystem(
        *System, *SourceEmitter, FGuid(), /*bCreateCopy=*/true);
    if (!HandleId.IsValid()) return FString();

    // Find the handle we just added by id and (optionally) rename it.
    for (FNiagaraEmitterHandle& H : System->GetEmitterHandles())
    {
        if (H.GetId() == HandleId)
        {
            if (!EmitterName.IsEmpty())
            {
                H.SetName(FName(*EmitterName), *System);
            }
            return H.GetName().ToString();
        }
    }
    return FString();
}

bool UMCPNiagaraLibrary::RemoveEmitterFromSystem(UNiagaraSystem* System, const FString& EmitterName)
{
    if (!System) return false;
    FNiagaraEmitterHandle* Handle = FindEmitterHandleMutable(System, EmitterName);
    if (!Handle) return false;
    System->Modify();
    System->RemoveEmitterHandle(*Handle);
    return true;
}

TArray<FString> UMCPNiagaraLibrary::ListModules(UNiagaraSystem* System, const FString& EmitterName, const FString& GroupName)
{
    TArray<FString> Out;
    const FStackLocation Loc = LocateStack(System, EmitterName, GroupName);
    if (!Loc.bValid()) return Out;

    TArray<UNiagaraNodeFunctionCall*> ModuleNodes;
    GatherFunctionCallNodes(Loc.OutputNode, ModuleNodes);
    for (UNiagaraNodeFunctionCall* N : ModuleNodes)
    {
        if (N) Out.Add(N->GetFunctionName());
    }
    return Out;
}

FString UMCPNiagaraLibrary::AddModuleToStack(UNiagaraSystem* System, const FString& EmitterName, const FString& GroupName, UNiagaraScript* ModuleScript, int32 Index)
{
    if (!ModuleScript) return FString();
    const FStackLocation Loc = LocateStack(System, EmitterName, GroupName);
    if (!Loc.bValid()) return FString();

    System->Modify();
    UNiagaraNodeFunctionCall* Added = FNiagaraStackGraphUtilities::AddScriptModuleToStack(
        ModuleScript, *Loc.OutputNode, Index, FString(), FGuid());
    if (!Added) return FString();
    return Added->GetFunctionName();
}

bool UMCPNiagaraLibrary::RemoveModuleFromStack(UNiagaraSystem* System, const FString& EmitterName, const FString& GroupName, const FString& ModuleName)
{
    const FStackLocation Loc = LocateStack(System, EmitterName, GroupName);
    if (!Loc.bValid()) return false;

    TArray<UNiagaraNodeFunctionCall*> ModuleNodes;
    GatherFunctionCallNodes(Loc.OutputNode, ModuleNodes);

    for (UNiagaraNodeFunctionCall* N : ModuleNodes)
    {
        if (N && N->GetFunctionName().Equals(ModuleName, ESearchCase::IgnoreCase))
        {
            System->Modify();
            UEdGraph* Graph = N->GetGraph();
            if (Graph)
            {
                Graph->Modify();
                N->BreakAllNodeLinks();
                Graph->RemoveNode(N);
            }
            return true;
        }
    }
    return false;
}

FString UMCPNiagaraLibrary::AddRendererToEmitter(UNiagaraSystem* System, const FString& EmitterName, UClass* RendererClass)
{
    if (!RendererClass || !RendererClass->IsChildOf(UNiagaraRendererProperties::StaticClass())) return FString();
    FNiagaraEmitterHandle* Handle = FindEmitterHandleMutable(System, EmitterName);
    if (!Handle) return FString();
    FVersionedNiagaraEmitter Versioned = Handle->GetInstance();
    UNiagaraEmitter* Emitter = Versioned.Emitter;
    if (!Emitter) return FString();

    Emitter->Modify();
    UNiagaraRendererProperties* Renderer = NewObject<UNiagaraRendererProperties>(Emitter, RendererClass, NAME_None, RF_Transactional);
    Emitter->AddRenderer(Renderer, Versioned.Version);
    return RendererClass->GetName();
}

int32 UMCPNiagaraLibrary::RemoveRenderersFromEmitter(UNiagaraSystem* System, const FString& EmitterName, UClass* RendererClass)
{
    if (!RendererClass) return 0;
    FNiagaraEmitterHandle* Handle = FindEmitterHandleMutable(System, EmitterName);
    if (!Handle) return 0;
    FVersionedNiagaraEmitter Versioned = Handle->GetInstance();
    UNiagaraEmitter* Emitter = Versioned.Emitter;
    FVersionedNiagaraEmitterData* Data = Versioned.GetEmitterData();
    if (!Emitter || !Data) return 0;

    Emitter->Modify();
    TArray<UNiagaraRendererProperties*> ToRemove;
    for (UNiagaraRendererProperties* R : Data->GetRenderers())
    {
        if (R && R->IsA(RendererClass)) ToRemove.Add(R);
    }
    for (UNiagaraRendererProperties* R : ToRemove)
    {
        Emitter->RemoveRenderer(R, Versioned.Version);
    }
    return ToRemove.Num();
}

bool UMCPNiagaraLibrary::RequestCompile(UNiagaraSystem* System)
{
    if (!System) return false;
    return System->RequestCompile(true);
}