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

#include "InworldBlueprintFunctionLibrary.h"

#include "Dom/JsonObject.h"
#include "Graph/InworldData.h"
#include "JsonObjectConverter.h"
#include "Policies/CondensedJsonPrintPolicy.h"
#include "Runtime/Launch/Resources/Version.h"
#include "UObject/TextProperty.h"
#if ENGINE_MAJOR_VERSION >= 5 && ENGINE_MINOR_VERSION >= 4
#include "Blueprint/BlueprintExceptionInfo.h"
#endif

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

#define LOCTEXT_NAMESPACE "InworldBlueprintFunctionLibrary"

void FJinjaArgument::Switch(TFunction<void(bool)> BoolFn, TFunction<void(float)> FloatFn, TFunction<void(int32)> IntFn,
	TFunction<void(const FString&)> StringFn, TFunction<void(TSharedPtr<FJsonObject>)> StructFn) const
{
	switch (ArgType)
	{
		case EJinjaArgumentType::BOOL:
			BoolFn(Bool);
			break;
		case EJinjaArgumentType::FLOAT:
			FloatFn(Float);
			break;
		case EJinjaArgumentType::INT:
			IntFn(Int);
			break;
		case EJinjaArgumentType::STRING:
			StringFn(String);
			break;
		case EJinjaArgumentType::STRUCT:
			StructFn(JsonObject.ToSharedPtr());
		default:
			break;
	}
}

void UInworldBlueprintFunctionLibrary::UnwrapInworldDataHandle(
	const FInworldDataHandle& DataHandle, UScriptStruct* ScriptStruct, int32& Data, bool& UnwrapSuccessful)
{
	// To Unwrap Data Handle in C++ directly, use Handle.Unwrap<Type>() to return underlying Type*.
	// Note: Only types derrived from FInworldData can be unwrapped.
	//		 Otherwise, type must be unwrapped to FInworldData_Struct, FInworldData_Object, etc.
	//		 Then, type can be accessed from bool FInworldData_Struct::GetData<Type>(OutType).
	checkNoEntry();
}

DEFINE_FUNCTION(UInworldBlueprintFunctionLibrary::execUnwrapInworldDataHandle)
{
	P_GET_STRUCT_REF(FInworldDataHandle, DataHandle)
	P_GET_OBJECT_REF(UScriptStruct, ScriptStruct);

	Stack.MostRecentPropertyAddress = nullptr;
	Stack.MostRecentPropertyContainer = nullptr;
	Stack.StepCompiledIn<FStructProperty>(nullptr);

	const FStructProperty* DataProp = CastField<FStructProperty>(Stack.MostRecentProperty);
	void* DataPtr = Stack.MostRecentPropertyAddress;

	P_GET_UBOOL_REF(UnwrapSuccessful);

	P_FINISH;

	UnwrapSuccessful = false;
	P_NATIVE_BEGIN;
	if (DataHandle.IsValid())
	{
		if ((*DataHandle)->IsA(ScriptStruct))
		{
			DataProp->Struct->CopyScriptStruct(DataPtr, *DataHandle);
			UnwrapSuccessful = true;
		}
		else if ((*DataHandle)->IsA(FInworldData_Struct::StaticStruct()))
		{
			FInworldData_Struct* Struct{DataHandle.Unwrap<FInworldData_Struct>()};
			if (Struct->GetDataScriptStruct() == ScriptStruct)
			{
				DataProp->Struct->CopyScriptStruct(DataPtr, Struct->GetData());
				UnwrapSuccessful = true;
			}
		}
	}
	P_NATIVE_END;
}

FInworldDataHandle UInworldBlueprintFunctionLibrary::WrapInworldData(UScriptStruct* ScriptStruct, const int32& Data)
{
	// To Wrap Data into Handle in C++ directly, use InworldData.Wrap() to return handle containing data.
	// Note: Only types derrived from FInworldData can be directly wrapped.
	//		 Otherwise, type must be contained inside FInworldData_Struct, FInworldData_Object, etc.
	checkNoEntry();
	return {};
}

DEFINE_FUNCTION(UInworldBlueprintFunctionLibrary::execWrapInworldData)
{
	P_GET_OBJECT_REF(UScriptStruct, ScriptStruct);

	Stack.MostRecentPropertyAddress = nullptr;
	Stack.MostRecentPropertyContainer = nullptr;
	Stack.StepCompiledIn<FStructProperty>(nullptr);

	void* ValuePtr = Stack.MostRecentPropertyAddress;

	P_FINISH;

	P_NATIVE_BEGIN;
	if (ScriptStruct->IsChildOf(FInworldData::StaticStruct()))
	{
		*StaticCast<FInworldDataHandle*>(RESULT_PARAM) = static_cast<FInworldData*>(ValuePtr)->Wrap();
	}
	else
	{
		*StaticCast<FInworldDataHandle*>(RESULT_PARAM) = FInworldData_Struct(ScriptStruct, ValuePtr);
	}
	P_NATIVE_END;
}

void UInworldBlueprintFunctionLibrary::CreateInworldDataError(
	UScriptStruct* InworldDataStructType, const FString& Reason, int32& Data)
{
	// To create Error from C++ directly, use FInworldData_Type::Error(FString Message).
	checkNoEntry();
}

DEFINE_FUNCTION(UInworldBlueprintFunctionLibrary::execCreateInworldDataError)
{
	P_GET_OBJECT_REF(UScriptStruct, InworldDataStructType);
	P_GET_PROPERTY_REF(FStrProperty, Reason);

	Stack.MostRecentPropertyAddress = nullptr;
	Stack.MostRecentPropertyContainer = nullptr;
	Stack.StepCompiledIn<FStructProperty>(nullptr);

	const FStructProperty* DataProp = CastField<FStructProperty>(Stack.MostRecentProperty);
	void* DataPtr = Stack.MostRecentPropertyAddress;

	P_FINISH;

	P_NATIVE_BEGIN;
	static_cast<FInworldData*>(DataPtr)->ErrorInfo = Reason;
	P_NATIVE_END;
}

void UInworldBlueprintFunctionLibrary::GetRuntimeStruct(
	UInworldProcessContext* ProcessContext, const FName& Key, UScriptStruct* ScriptStruct, int32& Struct, bool& Valid)
{
	// To get runtime data from C++ directly, use UInworldProcessContext functions to aquire InworldDataHandle, then unwrap.
	checkNoEntry();
}

DEFINE_FUNCTION(UInworldBlueprintFunctionLibrary::execGetRuntimeStruct)
{
	P_GET_OBJECT(UObject, ProcessContextObject)
	P_GET_PROPERTY_REF(FNameProperty, Key);
	P_GET_OBJECT_REF(UScriptStruct, ScriptStruct);

	Stack.MostRecentPropertyAddress = nullptr;
	Stack.MostRecentPropertyContainer = nullptr;
	Stack.StepCompiledIn<FStructProperty>(nullptr);

	const FStructProperty* DataProp = CastField<FStructProperty>(Stack.MostRecentProperty);
	void* DataPtr = Stack.MostRecentPropertyAddress;

	P_GET_UBOOL_REF(Valid);

	P_FINISH;

	Valid = false;
	P_NATIVE_BEGIN;
	if (UInworldProcessContext* ProcessContext = Cast<UInworldProcessContext>(ProcessContextObject))
	{
		if (FInworldDataHandle* DataHandle = ProcessContext->RuntimeDataMap.Find(Key); DataHandle != nullptr)
		{
			if ((**DataHandle)->IsA(ScriptStruct))
			{
				DataProp->Struct->CopyScriptStruct(DataPtr, **DataHandle);
				Valid = true;
			}
			else if ((**DataHandle)->IsA(FInworldData_Struct::StaticStruct()))
			{
				FInworldData_Struct* Struct{DataHandle->Unwrap<FInworldData_Struct>()};
				if (Struct->GetDataScriptStruct() == ScriptStruct)
				{
					DataProp->Struct->CopyScriptStruct(DataPtr, Struct->GetData());
					Valid = true;
				}
			}
		}
	}
	P_NATIVE_END;
}

void UInworldBlueprintFunctionLibrary::TrimLLMOutput(const FString& LLMOutputString, FString& OutTrimmedLLMOutputString)
{
	OutTrimmedLLMOutputString = LLMOutputString.Replace(TEXT("  "), TEXT(" "));
	OutTrimmedLLMOutputString.TrimStartAndEndInline();
	OutTrimmedLLMOutputString.TrimQuotesInline();
}

bool UInworldBlueprintFunctionLibrary::RenderJinjaWithJson(const FString& Template, const FString& JsonString, FString& OutJinja)
{
	if (inworld::StatusOr_string status_or_string{inworld::RenderJinja(TCHAR_TO_UTF8(*Template), TCHAR_TO_UTF8(*JsonString))};
		status_or_string.ok())
	{
		OutJinja = UTF8_TO_TCHAR(status_or_string.value().c_str());
		return true;
	}
	else
	{
		OutJinja = UTF8_TO_TCHAR(status_or_string.status().ToString().c_str());
		return false;
	}
}

FJinjaArgument UInworldBlueprintFunctionLibrary::CreateBooleanJinjaArgument(bool Value)
{
	return Value;
}

FJinjaArgument UInworldBlueprintFunctionLibrary::CreateFloatJinjaArgument(float Value)
{
	return Value;
}

FJinjaArgument UInworldBlueprintFunctionLibrary::CreateIntJinjaArgument(int32 Value)
{
	return Value;
}

FJinjaArgument UInworldBlueprintFunctionLibrary::CreateStringJinjaArgument(FString Value)
{
	return Value;
}

FJinjaArgument UInworldBlueprintFunctionLibrary::CreateStructJinjaArgument(const int32& Data)
{
	// To create FJinjaArgument from C++ directly, use ctor taking USTRUCT.
	checkNoEntry();
	return {};
}

DEFINE_FUNCTION(UInworldBlueprintFunctionLibrary::execCreateStructJinjaArgument)
{
	Stack.MostRecentPropertyAddress = nullptr;
	Stack.MostRecentPropertyContainer = nullptr;
	Stack.StepCompiledIn<FStructProperty>(nullptr);

	const FStructProperty* DataProp = CastField<FStructProperty>(Stack.MostRecentProperty);
	void* DataPtr = Stack.MostRecentPropertyAddress;

	P_FINISH;

	P_NATIVE_BEGIN;
	*StaticCast<FJinjaArgument*>(RESULT_PARAM) = FJinjaArgument(DataProp->Struct, DataPtr);
	P_NATIVE_END;
}

bool UInworldBlueprintFunctionLibrary::RenderJinjaWithArgMap(
	const FString& Template, const TMap<FString, FJinjaArgument>& ArgMap, FString& OutJinja)
{
	TSharedRef<FJsonObject> JsonObject = MakeShareable(new FJsonObject());

	for (const auto& Arg : ArgMap)
	{
		Arg.Value.Switch(
			[&](bool Value)
			{
				JsonObject->SetBoolField(Arg.Key, Value);
			},
			[&](float Value)
			{
				JsonObject->SetNumberField(Arg.Key, Value);
			},
			[&](int32 Value)
			{
				JsonObject->SetNumberField(Arg.Key, Value);
			},
			[&](const FString& Value)
			{
				JsonObject->SetStringField(Arg.Key, Value);
			},
			[&](TSharedPtr<FJsonObject> Value)
			{
				JsonObject->SetObjectField(Arg.Key, Value);
			});
	}

	FString JsonString;
	TSharedRef<TJsonWriter<TCHAR, TCondensedJsonPrintPolicy<TCHAR>>> JsonWriter =
		TJsonWriterFactory<TCHAR, TCondensedJsonPrintPolicy<TCHAR>>::Create(&JsonString, /*Indent=*/0);
	FJsonSerializer::Serialize(JsonObject, JsonWriter);
	return RenderJinjaWithJson(Template, JsonString, OutJinja);
}

FInworldVisemeBlends UInworldBlueprintFunctionLibrary::GetVisemeBlends(
	const float PlaybackTimeSeconds, const TArray<FInworldVisemeInfo>& VisemeInfos)
{
	FInworldVisemeBlends VisemeBlends = {};
	const int32 NumVisemes = VisemeInfos.Num();
	if (NumVisemes == 0)
	{
		return VisemeBlends;
	}
	// for initial viseme, blend stop with first viseme
	if (PlaybackTimeSeconds < VisemeInfos[0].Timestamp)
	{
		float Blend = PlaybackTimeSeconds / VisemeInfos[0].Timestamp;
		VisemeBlends.STOP = FMath::Clamp(1.f - Blend, 0.f, 1.f);
		VisemeBlends[VisemeInfos[0].Code] = FMath::Clamp(Blend, 0.f, 1.f);
		return VisemeBlends;
	}
	// for last visem
	if (PlaybackTimeSeconds > VisemeInfos.Last().Timestamp)
	{
		return VisemeBlends;
	}

	int32 TargetIndex = 0;
	for (int32 i = 1; i < NumVisemes; i++)
	{
		if (VisemeInfos[i].Timestamp >= PlaybackTimeSeconds)
		{
			TargetIndex = i;
			break;
		}
	}

	const FInworldVisemeInfo& CurrentVisemeInfo = VisemeInfos[TargetIndex];
	const FInworldVisemeInfo& PreviouseVisemeInfo = VisemeInfos[TargetIndex - 1];
	// if current viseme equal next
	if (CurrentVisemeInfo.Code == PreviouseVisemeInfo.Code)
	{
		VisemeBlends[CurrentVisemeInfo.Code] = 1.f;
		return VisemeBlends;
	}
	const float TimeDiff = CurrentVisemeInfo.Timestamp - PreviouseVisemeInfo.Timestamp;
	const float Blend = (TimeDiff > 0) ? (PlaybackTimeSeconds - PreviouseVisemeInfo.Timestamp) / TimeDiff : 0.f;
	VisemeBlends.STOP = 0.f;
	VisemeBlends[PreviouseVisemeInfo.Code] = FMath::Clamp(1.f - Blend, 0.f, 1.f);
	VisemeBlends[CurrentVisemeInfo.Code] = FMath::Clamp(Blend, 0.f, 1.f);
	return VisemeBlends;
}

FInworldDataHandle UInworldBlueprintFunctionLibrary::ErrorOnDestroying()
{
	return FInworldData_Error(FString::Printf(TEXT("UE Node object at destroying process")));
}

bool UInworldBlueprintFunctionLibrary::DataStreamTTSOutputHasNext(const FInworldData_DataStream_TTSOutput& TTSOutputStream)
{
	return TTSOutputStream.HasNext();
}

FInworldDataHandle UInworldBlueprintFunctionLibrary::DataStreamTTSOutputGetNext(FInworldData_DataStream_TTSOutput TTSOutputStream)
{
	return TTSOutputStream.GetNext();
}

#undef LOCTEXT_NAMESPACE
