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

#include "Graph/Nodes/InworldNode_TextClassifier.h"

#include "Dom/JsonObject.h"
#include "InworldRuntimeSettings.h"
#include "InworldRuntimeUtils.h"
#include "Serialization/JsonSerializer.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

UInworldNode_TextClassifier::UInworldNode_TextClassifier()
{
	NodeName = "Text Classifier";
}

EInworldRuntimeErrors UInworldNode_TextClassifier::GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson)
{
	const FString ComponentId = FString("Embedder_").Append(CreationConfig.EmbedderComponentId);

	TSharedPtr<FJsonObject> NodeObj = MakeShared<FJsonObject>();
	NodeObj->SetStringField(TEXT("id"), GetNodeName());
	NodeObj->SetStringField(TEXT("type"), TEXT("TextClassifierNode"));

	// execution_config
	{
		TSharedPtr<FJsonObject> ExecCfg = MakeShared<FJsonObject>();
		if (ExecutionConfig.HasTextClassifierConfig)
		{
			TSharedPtr<FJsonObject> ExecProps = MakeShared<FJsonObject>();
			// text_classifier_config
			{
				TSharedPtr<FJsonObject> ClassifierConfig = MakeShared<FJsonObject>();
				TArray<TSharedPtr<FJsonValue>> Classes;
				for (const FInworldClassThreshold& Class : ExecutionConfig.ClassifierConfig.Classes)
				{
					TSharedPtr<FJsonObject> ClassThreshold = MakeShared<FJsonObject>();
					ClassThreshold->SetStringField(TEXT("label"), Class.ClassLabel);
					ClassThreshold->SetNumberField(TEXT("threshold"), Class.Threshold);
					Classes.Add(MakeShared<FJsonValueObject>(ClassThreshold));
				}
				ClassifierConfig->SetArrayField(TEXT("classes"), Classes);
				ExecProps->SetObjectField(TEXT("classifier_config"), ClassifierConfig);
			}
			ExecCfg->SetObjectField(TEXT("properties"), ExecProps);
		}
		ExecCfg->SetStringField(TEXT("type"), TEXT("TextClassifierNodeExecutionConfig"));
		NodeObj->SetObjectField(TEXT("execution_config"), ExecCfg);
	}

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

	// creation_config
	{
		// text embedder component
		{
			TArray<TSharedPtr<FJsonValue>> Components = GraphJson->GetArrayField(TEXT("components"));
			bool bExists = false;
			for (const TSharedPtr<FJsonValue>& Component : Components)
			{
				const TSharedPtr<FJsonObject> ComponentObj = Component->AsObject();
				if (ComponentObj.IsValid() && ComponentObj->GetStringField(TEXT("id")) == ComponentId)
				{
					bExists = true;
					break;
				}
			}
			if (!bExists)
			{
				TSharedPtr<FJsonObject> Comp = MakeShared<FJsonObject>();
				Comp->SetStringField(TEXT("id"), ComponentId);
				Comp->SetStringField(TEXT("type"), TEXT("TextEmbedderInterface"));

				if (const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>())
				{
					if (const FInworldTextEmbedderCreationConfig* TextEmbedderCfg =
							Settings->TextEmbedderCreationConfig.Find(CreationConfig.EmbedderComponentId))
					{
						TSharedPtr<FJsonObject> CreationCfg = MakeShared<FJsonObject>();
						TSharedPtr<FJsonObject> CreationProps = MakeShared<FJsonObject>();

						if (TextEmbedderCfg->ComputeHost == EInworldComputeHost::REMOTE)
						{
							CreationCfg->SetStringField(TEXT("type"), TEXT("RemoteTextEmbedderConfig"));
							CreationProps->SetStringField(TEXT("provider"), TextEmbedderCfg->RemoteCreationConfig.Provider);
							CreationProps->SetStringField(TEXT("model_name"), TextEmbedderCfg->RemoteCreationConfig.Model);
							CreationProps->SetStringField(TEXT("api_key"), TEXT("{{INWORLD_API_KEY}}"));
						}
						else
						{
							CreationCfg->SetStringField(TEXT("type"), TEXT("LocalTextEmbedderConfig"));
							CreationProps->SetStringField(
								TEXT("model_path"), FPaths::Combine("{{INWORLD_EXTRAS_PATH}}", "embedders",
														TextEmbedderCfg->LocalCreationConfig.Model));
							TSharedPtr<FJsonObject> DeviceObj = MakeShared<FJsonObject>();
							inworld::Device device{FInworldRuntimeUtils::GetDevice(
								(TextEmbedderCfg->LocalCreationConfig.LocalComputeHost == EInworldLocalComputeHost::CPU)
									? inworld::DeviceType::kCPU
									: inworld::DeviceType::kCUDA)};
							DeviceObj->SetStringField(
								TEXT("type"), device.type() == inworld::DeviceType::kCPU ? TEXT("CPU") : TEXT("CUDA"));
							DeviceObj->SetNumberField(TEXT("index"), device.index());

							TSharedPtr<FJsonObject> DeviceInfoObj = MakeShared<FJsonObject>();
							DeviceInfoObj->SetStringField(TEXT("name"), UTF8_TO_TCHAR(device.info().name().c_str()));
							DeviceInfoObj->SetNumberField(TEXT("free_memory_bytes"), device.info().free_memory_bytes());
							DeviceInfoObj->SetNumberField(TEXT("total_memory_bytes"), device.info().total_memory_bytes());

							DeviceObj->SetObjectField(TEXT("info"), DeviceInfoObj);
							CreationProps->SetObjectField(TEXT("device"), DeviceObj);
						}

						CreationCfg->SetObjectField(TEXT("properties"), CreationProps);
						Comp->SetObjectField(TEXT("creation_config"), CreationCfg);
					}
				}
				Components.Add(MakeShared<FJsonValueObject>(Comp));
				GraphJson->SetArrayField(TEXT("components"), Components);
			}
		}

		TSharedPtr<FJsonObject> CreateCfg = MakeShared<FJsonObject>();
		{
			TSharedPtr<FJsonObject> CreateProps = MakeShared<FJsonObject>();
			CreateProps->SetStringField(TEXT("embedder_component_id"), ComponentId);
			// classifier_config
			{
				TSharedPtr<FJsonObject> ClassifierConfig = MakeShared<FJsonObject>();
				ClassifierConfig->SetStringField(TEXT("model_weights_path"),
					FPaths::Combine("{{INWORLD_EXTRAS_PATH}}", CreationConfig.ClassifierConfig.ModelWeightsPath));
				TArray<TSharedPtr<FJsonValue>> Classes;
				for (const FString& Class : CreationConfig.ClassifierConfig.SupportedClasses)
				{
					Classes.Add(MakeShared<FJsonValueString>(Class));
				}
				ClassifierConfig->SetArrayField(TEXT("supported_classes"), Classes);
				CreateProps->SetObjectField(TEXT("classifier_config"), ClassifierConfig);
			}
			CreateCfg->SetObjectField(TEXT("properties"), CreateProps);
		}
		CreateCfg->SetStringField(TEXT("type"), TEXT("TextClassifierNodeCreationConfig"));
		NodeObj->SetObjectField(TEXT("creation_config"), CreateCfg);
	}

	return EInworldRuntimeErrors::SUCCESS;
}
