/**
 * Copyright 2022-2025 Theai, Inc. dba Inworld AI
 */

#include "Graph/Edges/InworldEdge_WithCondition.h"

#include "Async/Async.h"
#include "GenericPlatform/GenericPlatformProcess.h"
#include "Graph/Assets/InworldGraphAsset.h"
#include "Graph/InworldData.h"
#include "Graph/InworldGraph.h"
#include "InworldRuntimeSubsystem.h"
#include "InworldRuntimeUtils.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

namespace inworld
{
bool ExecuteMeetsCondition(UInworldEdge_WithCondition* Edge, UInworldProcessContext* ProcessContext, FInworldDataHandle Input)
{
	bool Result{true};
	if (Edge->ShouldExecuteMeetsConditionInGameThread())
	{
		auto Event{FGenericPlatformProcess::GetSynchEventFromPool()};
		AsyncTask(ENamedThreads::GameThread,
			[ProcessContext, process_context = ProcessContext->GetProcessContext(), Input, Result = &Result, Event,
				WeakEdge = TWeakObjectPtr<UInworldEdge_WithCondition>(Edge)]()
			{
				ProcessContext->SetProcessContext(process_context);
				if (WeakEdge.IsValid())
				{
					*Result = WeakEdge->MeetsCondition(ProcessContext, Input) ^ WeakEdge->ShouldNegateCondition();
				}
				ProcessContext->SetProcessContext(nullptr);
				Event->Trigger();
			});
		Event->Wait();
		FGenericPlatformProcess::ReturnSynchEventToPool(Event);
	}
	else
	{
		Result = Edge->MeetsCondition(ProcessContext, Input) ^ Edge->ShouldNegateCondition();
	}
	return Result;
}
}	 // namespace inworld

TArray<UFunction*> UInworldEdge_WithCondition::GetMeetsConditionFunctions() const
{
	TArray<UFunction*> MeetsConditionFunctions;
	for (TFieldIterator<UFunction> FuncIt(GetClass(), EFieldIteratorFlags::IncludeSuper); FuncIt; ++FuncIt)
	{
		if (UFunction* const Function = *FuncIt; Function != nullptr && Function->GetName().StartsWith(TEXT("MeetsCondition")))
		{
			MeetsConditionFunctions.Add(Function);
		}
	}
	return MeetsConditionFunctions;
}

bool UInworldEdge_WithCondition::MeetsCondition(const UInworldProcessContext* ProcessContext, const FInworldDataHandle& Input) 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 : GetMeetsConditionFunctions())
	{
		bool bInvalid{false};
		bool bHasProcessContext{false};
		FProperty* ProcessContextFunctionProperty{nullptr};
		TOptional<FParamInfo> InputParamInfoOpt{};
		FProperty* ReturnValueFunctionProperty{nullptr};
		bool bHasBoolReturn{false};
		for (TFieldIterator<FProperty> PropIt(ProcessFunction); PropIt; ++PropIt)
		{
			if (FProperty* const Property = *PropIt; Property->HasAnyPropertyFlags(CPF_Parm))
			{
				const bool bIsReturn =
					Property->HasAnyPropertyFlags(CPF_ReturnParm) ||
					(Property->HasAnyPropertyFlags(CPF_OutParm) && !Property->HasAnyPropertyFlags(CPF_ConstParm));
				if (bIsReturn)
				{
					if (FBoolProperty* const BoolProperty = CastField<FBoolProperty>(Property); BoolProperty != nullptr)
					{
						bHasBoolReturn = true;
						ReturnValueFunctionProperty = Property;
					}
					else
					{
						// Return Value must be Bool.
						bInvalid = true;
						break;
					}
				}
				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 && !InputParamInfoOpt.IsSet())
							{
								bInvalid = true;
								break;
							}
							bHasProcessContext = true;
							ProcessContextFunctionProperty = Property;
							continue;
						}
					}

					// Must not exceed the 1 input argument
					if (InputParamInfoOpt.IsSet())
					{
						bInvalid = true;
						break;
					}

					FParamInfo* ParamInfo{nullptr};
					InputParamInfoOpt = FParamInfo{};
					ParamInfo = &InputParamInfoOpt.GetValue();
					ParamInfo->Size = Property->GetSize();
					ParamInfo->FunctionProperty = Property;

					if (FStructProperty* const StructProperty = CastField<FStructProperty>(Property); StructProperty != nullptr)
					{
						// Input must be:
						//		FInworldDataHandle
						//		FInworldData_Type, where Type matches
						//		FInworldData_Struct, where underlying ScriptStruct matches
						if (!(StructProperty->Struct == FInworldDataHandle::StaticStruct() || Input->IsA(StructProperty->Struct) ||
								(Input->IsA(FInworldData_Struct::StaticStruct()) &&
									Input.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 && Input.Unwrap<FInworldData_Struct>() &&
							Input.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 && Input.Unwrap<FInworldData_Struct>() &&
							Input.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 && Input.Unwrap<FInworldData_Struct>() &&
							Input.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 && Input.Unwrap<FInworldData_Struct>() &&
							Input.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;
					}
				}
			}
		}

		if (bInvalid || !bHasBoolReturn || !InputParamInfoOpt.IsSet())
		{
			continue;
		}

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

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

		{
			const FParamInfo& InputParamInfo{InputParamInfoOpt.GetValue()};
			void* Write = InputParamInfo.FunctionProperty->ContainerPtrToValuePtr<void*>(Params);
			if (InputParamInfo.Type == EParamType::HANDLE)
			{
				FMemory::Memcpy(Write, (void*) &Input, InputParamInfo.Size);
			}
			else if (InputParamInfo.Type == EParamType::DATA)
			{
				FMemory::Memcpy(Write, (void*) Input.Data.Get(), InputParamInfo.Size);
			}
			else if (InputParamInfo.Type == EParamType::CUSTOM)
			{
				FMemory::Memcpy(Write, ((FInworldData_Struct*) Input.Data.Get())->GetData(), InputParamInfo.Size);
			}
			else if (InputParamInfo.Type == EParamType::OBJECT)
			{
				FInworldData_Object Object;
				((FInworldData_Struct*) Input.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*) Input.Data.Get())->GetData(Bool);
				*InputParamInfo.FunctionProperty->ContainerPtrToValuePtr<bool>(Params) = Bool.Bool;
			}
			else if (InputParamInfo.Type == EParamType::INT)
			{
				FInworldData_Int Int;
				((FInworldData_Struct*) Input.Data.Get())->GetData(Int);
				*InputParamInfo.FunctionProperty->ContainerPtrToValuePtr<int32>(Params) = Int.Int;
			}
			else if (InputParamInfo.Type == EParamType::FLOAT)
			{
				FInworldData_Float Float;
				((FInworldData_Struct*) Input.Data.Get())->GetData(Float);
				*InputParamInfo.FunctionProperty->ContainerPtrToValuePtr<double>(Params) = Float.Float;
			}
		}

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

		void* Write = ReturnValueFunctionProperty->ContainerPtrToValuePtr<void*>(Params);
		bool Result = *(bool*) Write;
		FMemory::Free(Params);
		return Result;
	}

	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();
			}
		}
	}

	const FString EdgeId = FString::Printf(TEXT("%s_to_%s"), *Source, *Destination).Replace(TEXT(" "), TEXT(""));
	const FString ErrorReason{FString::Format(TEXT("Edge: {0} ({1}) - No valid 'MeetsCondition' function for input type:{1}"),
		{*EdgeId, *GetClass()->GetName(), *InputStruct->GetName()})};
	UE_LOG(LogInworldGraph, Error, TEXT("%s"), *ErrorReason);
	return false;
}

EInworldRuntimeErrors UInworldEdge_WithCondition::GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson)
{
	TSharedPtr<FJsonObject> Main = GraphJson->GetObjectField(TEXT("main"));
	const FString GraphId = Main->GetStringField(TEXT("id"));

	const FString ConditionId = FString::Printf(TEXT("%s/%s_to_%s"), *GraphId, *Source, *Destination).Replace(TEXT(" "), TEXT(""));

	TArray<TSharedPtr<FJsonValue>> Components = GraphJson->GetArrayField(TEXT("components"));
	{
		TSharedPtr<FJsonObject> ComponentObj = MakeShared<FJsonObject>();
		ComponentObj->SetStringField(TEXT("id"), ConditionId);
		ComponentObj->SetStringField(TEXT("type"), FString::Printf(TEXT("UCE/%s%s"), *ConditionId, *GetClass()->GetPathName()));
		Components.Add(MakeShared<FJsonValueObject>(ComponentObj));

		// create 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);
			ComponentObj->SetObjectField(TEXT("creation_config"), InstanceCfg);
		}

		GraphJson->SetArrayField(TEXT("components"), Components);
	}

	TSharedPtr<FJsonObject> EdgeObj = MakeShared<FJsonObject>();
	EdgeObj->SetStringField(TEXT("from_node"), Source);
	EdgeObj->SetStringField(TEXT("to_node"), Destination);
	EdgeObj->SetBoolField(TEXT("optional"), !IsRequired);
	EdgeObj->SetBoolField(TEXT("loop"), IsLoopEdge);
	EdgeObj->SetStringField(TEXT("condition_id"), ConditionId);

	TArray<TSharedPtr<FJsonValue>> Edges = Main->GetArrayField(TEXT("edges"));
	Edges.Add(MakeShared<FJsonValueObject>(EdgeObj));
	Main->SetArrayField(TEXT("edges"), Edges);

	return EInworldRuntimeErrors::SUCCESS;
}
