/**
 * Copyright 2022-2025 Theai, Inc. dba Inworld AI
 */
#include "Graph/Nodes/InworldNode_Custom.h"

#include "Async/Async.h"
#include "Engine/Engine.h"
#include "GenericPlatform/GenericPlatformMemory.h"
#include "GenericPlatform/GenericPlatformProcess.h"
#include "Graph/InworldData.h"
#include "Graph/InworldGraph.h"
#include "InworldBlueprintFunctionLibrary.h"
#include "InworldRuntimeSubsystem.h"
#include "InworldRuntimeUtils.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

namespace inworld
{
FInworldDataHandle ExecuteCustomNode(
	UInworldNode_Custom* Node, UInworldProcessContext* ProcessContext, const TArray<FInworldDataHandle>& Inputs)
{
	auto* Subsystem = UInworldRuntimeSubsystem::Get();
	if (!Node || !Subsystem || !Subsystem->IsNodeValid(Node))
	{
		return UInworldBlueprintFunctionLibrary::ErrorOnDestroying();
	}

	FInworldDataHandle Result;
	if (Node->ShouldExecuteProcessInGameThread())
	{
		auto Event{FGenericPlatformProcess::GetSynchEventFromPool()};
		AsyncTask(ENamedThreads::GameThread,
			[ProcessContext, process_context = ProcessContext->GetProcessContext(), Inputs, &Result, Event,
				WeakNode = TWeakObjectPtr<UInworldNode_Custom>(Node)]()
			{
				ProcessContext->SetProcessContext(process_context);
				if (WeakNode.IsValid())
				{
					Result = WeakNode->Process(ProcessContext, Inputs);
				}
				ProcessContext->SetProcessContext(nullptr);
				Event->Trigger();
			});
		Event->Wait();
		FGenericPlatformProcess::ReturnSynchEventToPool(Event);
	}
	else
	{
		Result = Node->Process(ProcessContext, Inputs);
	}

	Subsystem = UInworldRuntimeSubsystem::Get();
	if (!Node || !Subsystem || !Subsystem->IsNodeValid(Node))
	{
		Result = UInworldBlueprintFunctionLibrary::ErrorOnDestroying();
	}
	else if (!Result.IsValid())
	{
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldNode_Custom::Process() - %s returned empty data"), *Node->GetNodeName());
		Result = FInworldData_Error(FString::Printf(TEXT("%s returned empty data"), *Node->GetNodeName()));
	}
	return Result;
}

}	 // namespace inworld

FInworldDataHandle UInworldNode_Custom::Process(
	const UInworldProcessContext* ProcessContext, const TArray<struct FInworldDataHandle>& Inputs) const
{
	enum EParamType
	{
		UNKNOWN,
		HANDLE,
		DATA,
		CUSTOM,
		OBJECT,
		BOOL,
		INT,
		FLOAT,
	};

	struct FParamInfo
	{
		int32 Size{0};
		EParamType Type{EParamType::UNKNOWN};
		UScriptStruct* ScriptStruct{nullptr};
		FProperty* FunctionProperty{nullptr};
	};

	for (UFunction* const ProcessFunction : GetProcessFunctions())
	{
		bool bInvalid{false};
		bool bHasProcessContext{false};
		FProperty* ProcessContextFunctionProperty{nullptr};
		TArray<FParamInfo> InputParamInfos{};
		TOptional<FParamInfo> OutputParamInfoOpt{};
		for (TFieldIterator<FProperty> PropIt(ProcessFunction); PropIt; ++PropIt)
		{
			if (FProperty* Property = *PropIt; Property->HasAnyPropertyFlags(CPF_Parm))
			{
				const int NextInputsIndex{InputParamInfos.Num()};

				FParamInfo* ParamInfo{nullptr};
				const bool bIsReturn =
					Property->HasAnyPropertyFlags(CPF_ReturnParm) ||
					(Property->HasAnyPropertyFlags(CPF_OutParm) && !Property->HasAnyPropertyFlags(CPF_ConstParm));
				if (bIsReturn)
				{
					OutputParamInfoOpt = FParamInfo{};
					ParamInfo = &OutputParamInfoOpt.GetValue();
				}
				else
				{
					if (FObjectPropertyBase* const ObjectProperty = CastField<FObjectProperty>(Property); ObjectProperty != nullptr)
					{
						if (ObjectProperty->PropertyClass == UInworldProcessContext::StaticClass())
						{
							// It must be the first argument (if present), and only 1 may exist.
							if (bHasProcessContext || !InputParamInfos.IsEmpty())
							{
								bInvalid = true;
								break;
							}
							bHasProcessContext = true;
							ProcessContextFunctionProperty = Property;
							continue;
						}
					}

					// Must not exceed the number of input arguments
					if (NextInputsIndex == Inputs.Num())
					{
						bInvalid = true;
						break;
					}

					ParamInfo = &InputParamInfos.AddDefaulted_GetRef();
				}
				ParamInfo->Size = Property->GetSize();
				ParamInfo->FunctionProperty = Property;

				if (FStructProperty* const StructProperty = CastField<FStructProperty>(Property); StructProperty != nullptr)
				{
					if (!bIsReturn)
					{
						// Input must be:
						//		FInworldDataHandle
						//		FInworldData_Type, where Type matches
						//		FInworldData_Struct, where underlying ScriptStruct matches
						if (!(StructProperty->Struct == FInworldDataHandle::StaticStruct() ||
								Inputs[NextInputsIndex]->IsA(StructProperty->Struct) ||
								(Inputs[NextInputsIndex]->IsA(FInworldData_Struct::StaticStruct()) &&
									Inputs[NextInputsIndex].Unwrap<FInworldData_Struct>()->GetDataScriptStruct() ==
										StructProperty->Struct)))
						{
							bInvalid = true;
							break;
						}
					}

					ParamInfo->ScriptStruct = StructProperty->Struct;
					if (StructProperty->Struct == FInworldDataHandle::StaticStruct())
					{
						ParamInfo->Type = EParamType::HANDLE;
					}
					else if (StructProperty->Struct->IsChildOf(FInworldData::StaticStruct()))
					{
						ParamInfo->Type = EParamType::DATA;
					}
					else
					{
						ParamInfo->Type = EParamType::CUSTOM;
					}
				}
				else if (FObjectPropertyBase* const ObjectProperty = CastField<FObjectProperty>(Property);
						 ObjectProperty != nullptr)
				{
					if (!bIsReturn && Inputs[NextInputsIndex].Unwrap<FInworldData_Struct>() &&
						Inputs[NextInputsIndex].Unwrap<FInworldData_Struct>()->GetDataScriptStruct() !=
							FInworldData_Object::StaticStruct())
					{
						bInvalid = true;
						break;
					}
					ParamInfo->Type = EParamType::OBJECT;
				}
				else if (FBoolProperty* const BoolProperty = CastField<FBoolProperty>(Property); BoolProperty != nullptr)
				{
					if (!bIsReturn && Inputs[NextInputsIndex].Unwrap<FInworldData_Struct>() &&
						Inputs[NextInputsIndex].Unwrap<FInworldData_Struct>()->GetDataScriptStruct() !=
							FInworldData_Bool::StaticStruct())
					{
						bInvalid = true;
						break;
					}
					ParamInfo->Type = EParamType::BOOL;
				}
				else if (FIntProperty* const IntProperty = CastField<FIntProperty>(Property); IntProperty != nullptr)
				{
					if (!bIsReturn && Inputs[NextInputsIndex].Unwrap<FInworldData_Struct>() &&
						Inputs[NextInputsIndex].Unwrap<FInworldData_Struct>()->GetDataScriptStruct() !=
							FInworldData_Int::StaticStruct())
					{
						bInvalid = true;
						break;
					}
					ParamInfo->Type = EParamType::INT;
				}
				else if (FDoubleProperty* const FloatProperty = CastField<FDoubleProperty>(Property); FloatProperty != nullptr)
				{
					if (!bIsReturn && Inputs[NextInputsIndex].Unwrap<FInworldData_Struct>() &&
						Inputs[NextInputsIndex].Unwrap<FInworldData_Struct>()->GetDataScriptStruct() !=
							FInworldData_Float::StaticStruct())
					{
						bInvalid = true;
						break;
					}
					ParamInfo->Type = EParamType::FLOAT;
				}
				else
				{
					// Only UObject (UInworldProcessContext, UObject), UStruct (FInworldData/Handle), and (Bool/Int/Float) allowed.
					bInvalid = true;
					break;
				}
			}
		}

		// TODO: Graph shouldn't be attempting to send Input for this type of node...
		const bool bContextOnlyNode =
			bHasProcessContext && ProcessFunction->NumParms == 2 && OutputParamInfoOpt.IsSet() && InputParamInfos.IsEmpty();
		const bool bPureGetterNode =
			bIsGetter && ProcessFunction->NumParms == 1 && OutputParamInfoOpt.IsSet() && InputParamInfos.IsEmpty();

		if (bInvalid || (InputParamInfos.Num() != Inputs.Num() && !bContextOnlyNode && !bPureGetterNode) ||
			!OutputParamInfoOpt.IsSet())
		{
			continue;
		}

		void* Params = FMemory::Malloc(ProcessFunction->ParmsSize);

		if (bHasProcessContext)
		{
			const uint64 ProcessContextSize{sizeof(UInworldProcessContext*)};
			FMemory::Memcpy(
				ProcessContextFunctionProperty->ContainerPtrToValuePtr<void*>(Params), (void*) &ProcessContext, ProcessContextSize);
		}

		for (int32 InputParamIndex = 0; InputParamIndex < InputParamInfos.Num(); ++InputParamIndex)
		{
			const FParamInfo& InputParamInfo{InputParamInfos[InputParamIndex]};
			void* Write = InputParamInfo.FunctionProperty->ContainerPtrToValuePtr<void*>(Params);
			if (InputParamInfo.Type == EParamType::HANDLE)
			{
				FMemory::Memcpy(Write, (void*) &Inputs[InputParamIndex], InputParamInfo.Size);
			}
			else if (InputParamInfo.Type == EParamType::DATA)
			{
				FMemory::Memcpy(Write, (void*) Inputs[InputParamIndex].Data.Get(), InputParamInfo.Size);
			}
			else if (InputParamInfo.Type == EParamType::CUSTOM)
			{
				FMemory::Memcpy(Write, ((FInworldData_Struct*) Inputs[InputParamIndex].Data.Get())->GetData(), InputParamInfo.Size);
			}
			else if (InputParamInfo.Type == EParamType::OBJECT)
			{
				FInworldData_Object Object;
				((FInworldData_Struct*) Inputs[InputParamIndex].Data.Get())->GetData(Object);
				UObject* ObjectPtr{Object.Object};
				FMemory::Memcpy(Write, (void*) &ObjectPtr, InputParamInfo.Size);
			}
			else if (InputParamInfo.Type == EParamType::BOOL)
			{
				FInworldData_Bool Bool;
				((FInworldData_Struct*) Inputs[InputParamIndex].Data.Get())->GetData(Bool);
				*InputParamInfo.FunctionProperty->ContainerPtrToValuePtr<bool>(Params) = Bool.Bool;
			}
			else if (InputParamInfo.Type == EParamType::INT)
			{
				FInworldData_Int Int;
				((FInworldData_Struct*) Inputs[InputParamIndex].Data.Get())->GetData(Int);
				*InputParamInfo.FunctionProperty->ContainerPtrToValuePtr<int32>(Params) = Int.Int;
			}
			else if (InputParamInfo.Type == EParamType::FLOAT)
			{
				FInworldData_Float Float;
				((FInworldData_Struct*) Inputs[InputParamIndex].Data.Get())->GetData(Float);
				*InputParamInfo.FunctionProperty->ContainerPtrToValuePtr<double>(Params) = Float.Float;
			}
		}
		const FParamInfo& OutputParamInfo = OutputParamInfoOpt.GetValue();
		void* Write = OutputParamInfo.FunctionProperty->ContainerPtrToValuePtr<void*>(Params);
		if (OutputParamInfo.ScriptStruct)
		{
			OutputParamInfo.ScriptStruct->InitializeDefaultValue((uint8*) Write);
		}

		const_cast<UInworldNode_Custom*>(this)->ProcessEvent(ProcessFunction, Params);

		FInworldDataHandle Result{};
		if (OutputParamInfo.Type == EParamType::HANDLE)
		{
			Result = {*(FInworldDataHandle*) (Write)};
		}
		else if (OutputParamInfo.Type == EParamType::DATA)
		{
			Result = ((FInworldData*) (Write))->Wrap();
		}
		else if (OutputParamInfo.Type == EParamType::CUSTOM)
		{
			Result = FInworldData_Struct(OutputParamInfo.ScriptStruct, Write);
		}
		else if (OutputParamInfo.Type == EParamType::OBJECT)
		{
			Result = FInworldData_Object(*(TObjectPtr<UObject>*) Write);
		}
		else if (OutputParamInfo.Type == EParamType::BOOL)
		{
			Result = FInworldData_Bool(*(bool*) Write);
		}
		else if (OutputParamInfo.Type == EParamType::INT)
		{
			Result = FInworldData_Int(*(int32*) Write);
		}
		else if (OutputParamInfo.Type == EParamType::FLOAT)
		{
			Result = FInworldData_Float(*(double*) Write);
		}

		if (OutputParamInfo.ScriptStruct)
		{
			OutputParamInfo.ScriptStruct->DestroyStruct(Write);
		}

		FMemory::Free(Params);
		return Result;
	}

	const FString FormatSeparator{TEXT("\n\t")};
	TArray<FString> InputTypeNames{{}};
	for (const struct FInworldDataHandle& Input : Inputs)
	{
		UStruct* InputStruct = Input->GetScriptStruct();
		if (FInworldData_Struct* Struct = Input.Unwrap<FInworldData_Struct>(); Struct != nullptr)
		{
			InputStruct = Struct->GetDataScriptStruct();
			if (InputStruct == FInworldData_Object::StaticStruct())
			{
				if (FInworldData_Object Object; Struct->GetData(Object) && Object.Object != nullptr)
				{
					InputStruct = Object.Object->GetClass();
				}
			}
		}
		InputTypeNames.Push(InputStruct->GetName());
	}

	const FString ErrorReason{FString::Format(TEXT("{0} - No valid 'Process' function for input type(s):{1}"),
		{*GetClass()->GetName(), *FString::Join(InputTypeNames, *FormatSeparator)})};
	UE_LOG(LogInworldGraph, Error, TEXT("%s"), *ErrorReason);
	return FInworldData_Error::Error(ErrorReason);
}

EInworldRuntimeErrors UInworldNode_Custom::GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson)
{
	TSharedPtr<FJsonObject> NodeObj = MakeShared<FJsonObject>();
	NodeObj->SetStringField(TEXT("id"), GetNodeName());
	NodeObj->SetStringField(TEXT("type"), FString::Printf(TEXT("UCN%s"), *GetClass()->GetPathName()));

	// create instance config
	{
		TSharedPtr<FJsonObject> InstanceCfg = MakeShared<FJsonObject>();
		InstanceCfg->SetStringField(TEXT("type"), TEXT("UnrealObjectInstance"));
		TSharedPtr<FJsonObject> InstanceProps = MakeShared<FJsonObject>();
		InstanceProps->SetStringField(TEXT("InstanceData"), FInworldRuntimeUtils::SerializeInstance(this));
		InstanceCfg->SetObjectField(TEXT("properties"), InstanceProps);
		NodeObj->SetObjectField(TEXT("creation_config"), InstanceCfg);
	}

	// append to main.nodes
	TSharedPtr<FJsonObject> Main = GraphJson->GetObjectField(TEXT("main"));
	TArray<TSharedPtr<FJsonValue>> Nodes = Main->GetArrayField(TEXT("nodes"));
	Nodes.Add(MakeShared<FJsonValueObject>(NodeObj));
	Main->SetArrayField(TEXT("nodes"), Nodes);

	return EInworldRuntimeErrors::SUCCESS;
}
