// Copyright (c) 2026 Codemongous LLC. All Rights Reserved.

#include "SCW_UStateTree.h"
#if ISCW_UE_VERSION_NEWER_THAN_OR_EQUAL(5,6,0)
#include "SCW_FInstancedPropertyBag.h"
#include "StateTree.h"
#include "StateTreeEditorData.h"
#include "StateTreeEditingSubsystem.h"
#include "StateTreeCompilerLog.h"
#include "StateTreeState.h"
#include "StateTreeEditorPropertyBindings.h"
#include "StructUtils/PropertyBag.h"
#include "StructUtils/InstancedStruct.h"
#include "SourcifyCoreUEWrapperModule.h"

using namespace Sourcify;

SCW_UStateTree::SCW_UStateTree(UStateTree* InStateTree)
	: SCW_UDataAsset(InStateTree), StateTree(InStateTree)
{
}

UStateTree* SCW_UStateTree::ToUStateTreeRawPtr() const
{
	return StateTree;
}

ISCW_UStateTreePtr SCW_UStateTree::ToUStateTreePtr()
{
	return GetSharedThisAs<SCW_UStateTree>();
}

ISCW_UStateTreeConstPtr SCW_UStateTree::ToUStateTreePtr() const
{
	return GetSharedThisAs<SCW_UStateTree>();
}

void SCW_UStateTree::Compile()
{
	if (StateTree)
	{
		FStateTreeCompilerLog Log;
		UStateTreeEditingSubsystem::CompileStateTree(StateTree, Log);
	}
}

void SCW_UStateTree::VisitPropertyBags(function<void(ISCW_FInstancedPropertyBag&)> Visitor)
{
	if (!StateTree)
	{
		return;
	}

	UStateTreeEditorData* EditorData = Cast<UStateTreeEditorData>(StateTree->EditorData);
	if (!EditorData)
	{
		return;
	}

	// Visit root-level parameter property bag
	// Use const_cast since GetRootParametersPropertyBag returns const& but we need to mutate
	FInstancedPropertyBag& RootBag = const_cast<FInstancedPropertyBag&>(EditorData->GetRootParametersPropertyBag());
	if (RootBag.IsValid())
	{
		SCW_FInstancedPropertyBag Handle(RootBag);
		Visitor(Handle);
	}

	EditorData->VisitHierarchy([&Visitor](UStateTreeState& State, UStateTreeState* /*ParentState*/) -> EStateTreeVisitor
	{
		if (State.Parameters.Parameters.IsValid())
		{
			SCW_FInstancedPropertyBag Handle(State.Parameters.Parameters);
			Visitor(Handle);
		}
		return EStateTreeVisitor::Continue;
	});
}

void SCW_UStateTree::VisitSingleEditorNode(FStateTreeEditorNode& Node, function<void(ISCW_UScriptStructConstPtr, function<void(ISCW_UScriptStructPtr)>)>& Visitor)
{
	auto VisitInstancedStruct = [&](FInstancedStruct& Struct)
	{
		const UScriptStruct* CurrentStruct = Struct.GetScriptStruct();
		if (!CurrentStruct)
		{
			return;
		}

		ISCW_UScriptStructConstPtr WrappedCurrent = FSourcifyCoreUEWrapper::Get().Wrap(*const_cast<UScriptStruct*>(CurrentStruct));
		auto Replacer = [&Struct](ISCW_UScriptStructPtr NewType)
		{
			if (NewType)
			{
				UScriptStruct* RawNewStruct = Cast<UScriptStruct>(const_cast<UObject*>(FSourcifyCoreUEWrapper::Unwrap(NewType)));
				if (RawNewStruct)
				{
					Struct.InitializeAs(RawNewStruct);
				}
			}
		};
		Visitor(WrappedCurrent, Replacer);
	};

	VisitInstancedStruct(Node.Node);
	VisitInstancedStruct(Node.Instance);
#if ISCW_UE_VERSION_NEWER_THAN_OR_EQUAL(5,7,0)
	VisitInstancedStruct(Node.ExecutionRuntimeData);
#endif
}

void SCW_UStateTree::VisitEditorNodes(TArray<FStateTreeEditorNode>& Nodes, function<void(ISCW_UScriptStructConstPtr, function<void(ISCW_UScriptStructPtr)>)>& Visitor)
{
	for (FStateTreeEditorNode& Node : Nodes)
	{
		VisitSingleEditorNode(Node, Visitor);
	}
}

void SCW_UStateTree::VisitEditorNodeStructTypes(function<void(ISCW_UScriptStructConstPtr currentType, function<void(ISCW_UScriptStructPtr newType)> replacer)> Visitor)
{
	if (!StateTree)
	{
		return;
	}

	UStateTreeEditorData* EditorData = Cast<UStateTreeEditorData>(StateTree->EditorData);
	if (!EditorData)
	{
		return;
	}

	VisitEditorNodes(EditorData->Evaluators, Visitor);
	VisitEditorNodes(EditorData->GlobalTasks, Visitor);

	EditorData->VisitHierarchy([this, &Visitor](UStateTreeState& State, UStateTreeState* /*ParentState*/) -> EStateTreeVisitor
	{
		VisitEditorNodes(State.EnterConditions, Visitor);
		VisitEditorNodes(State.Tasks, Visitor);
		VisitEditorNodes(State.Considerations, Visitor);
		VisitSingleEditorNode(State.SingleTask, Visitor);

		for (FStateTreeTransition& Transition : State.Transitions)
		{
			VisitEditorNodes(Transition.Conditions, Visitor);
		}

		return EStateTreeVisitor::Continue;
	});
}

void SCW_UStateTree::VisitEventPayloadStructs(function<void(ISCW_UScriptStructConstPtr currentType, function<void(ISCW_UScriptStructPtr newType)> replacer)> Visitor)
{
	if (!StateTree)
	{
		return;
	}

	UStateTreeEditorData* EditorData = Cast<UStateTreeEditorData>(StateTree->EditorData);
	if (!EditorData)
	{
		return;
	}

	EditorData->VisitHierarchy([&Visitor](UStateTreeState& State, UStateTreeState* /*ParentState*/) -> EStateTreeVisitor
	{
		if (State.RequiredEventToEnter.PayloadStruct)
		{
			const UScriptStruct* CurrentStruct = State.RequiredEventToEnter.PayloadStruct;
			ISCW_UScriptStructConstPtr WrappedCurrent = FSourcifyCoreUEWrapper::Get().Wrap(*const_cast<UScriptStruct*>(CurrentStruct));
			auto Replacer = [&State](ISCW_UScriptStructPtr NewType)
			{
				if (NewType)
				{
					UScriptStruct* RawNewStruct = Cast<UScriptStruct>(const_cast<UObject*>(FSourcifyCoreUEWrapper::Unwrap(NewType)));
					State.RequiredEventToEnter.PayloadStruct = RawNewStruct;
				}
			};
			Visitor(WrappedCurrent, Replacer);
		}
		return EStateTreeVisitor::Continue;
	});
}

void SCW_UStateTree::VisitSchemaTypeReferences(function<void(ISCW_UObjectConstPtr currentType, function<void(ISCW_UObjectPtr newType)> replacer)> Visitor)
{
	if (!StateTree)
	{
		return;
	}

	UStateTreeEditorData* EditorData = Cast<UStateTreeEditorData>(StateTree->EditorData);
	if (!EditorData || !EditorData->Schema)
	{
		return;
	}

	UStateTreeSchema* Schema = EditorData->Schema;

	for (TFieldIterator<FObjectPropertyBase> It(Schema->GetClass()); It; ++It)
	{
		FObjectPropertyBase* ObjProp = *It;
		void* ValuePtr = ObjProp->ContainerPtrToValuePtr<void>(Schema);
		UObject* CurrentValue = ObjProp->GetObjectPropertyValue(ValuePtr);
		if (!CurrentValue)
		{
			continue;
		}

		ISCW_UObjectConstPtr WrappedCurrent = FSourcifyCoreUEWrapper::Get().Wrap(*CurrentValue);
		auto Replacer = [ObjProp, ValuePtr, Schema](ISCW_UObjectPtr NewType)
		{
			if (NewType)
			{
				UObject* RawNew = const_cast<UObject*>(FSourcifyCoreUEWrapper::Unwrap(NewType));
				ObjProp->SetObjectPropertyValue(ValuePtr, RawNew);

				// Fire PostEditChangeChainProperty (NOT PostEditChangeProperty) because
				// schema subclasses like UStateTreeComponentSchema override the chain
				// variant to sync ContextDataDescs from their source properties.
				// PostEditChangeProperty does NOT call PostEditChangeChainProperty.
				FEditPropertyChain PropertyChain;
				PropertyChain.AddHead(ObjProp);
				FPropertyChangedEvent PropertyChangedEvent(ObjProp, EPropertyChangeType::ValueSet);
				FPropertyChangedChainEvent ChainEvent(PropertyChain, PropertyChangedEvent);
				Schema->PostEditChangeChainProperty(ChainEvent);
			}
		};
		Visitor(WrappedCurrent, Replacer);
	}
}

void SCW_UStateTree::VisitPropertyBindingPathSegments(function<wstring(const wstring& segmentName)> Visitor)
{
	if (!StateTree)
	{
		return;
	}

	UStateTreeEditorData* EditorData = Cast<UStateTreeEditorData>(StateTree->EditorData);
	if (!EditorData)
	{
		return;
	}

	FStateTreeEditorPropertyBindings* EditorBindings = EditorData->GetPropertyEditorBindings();
	if (!EditorBindings)
	{
		return;
	}

	auto RenameSegments = [&Visitor](FPropertyBindingPath& Path)
	{
		for (FPropertyBindingPathSegment& Segment : Path.GetMutableSegments())
		{
			FName CurrentName = Segment.GetName();
			wstring CurrentNameStr = *CurrentName.ToString();

			wstring NewNameStr = Visitor(CurrentNameStr);
			if (!NewNameStr.empty() && NewNameStr != CurrentNameStr)
			{
				Segment.SetName(FName(NewNameStr.c_str()));
			}
		}
	};

	for (FStateTreePropertyPathBinding& Binding : EditorBindings->GetMutableBindings())
	{
		RenameSegments(Binding.GetMutableSourcePath());
		RenameSegments(Binding.GetMutableTargetPath());
	}
}
#endif // ISCW_UE_VERSION_NEWER_THAN_OR_EQUAL(5,6,0)