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

#include "Graph/Nodes/InworldNode_STT.h"

#include "Dom/JsonObject.h"
#include "Graph/InworldGraph.h"
#include "InworldRuntimeSettings.h"
#include "Misc/Paths.h"
#include "Serialization/JsonSerializer.h"

UInworldNode_STT::UInworldNode_STT()
{
	NodeName = "STT";
}

EInworldRuntimeErrors UInworldNode_STT::GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson)
{
	const FString ComponentId =
		ExecutionConfig.bStreamed ? ExecutionConfig.StreamingSTTComponentId : ExecutionConfig.STTComponentId;
	const FString STTComponentId =
		ExecutionConfig.bStreamed ? FString("StreamingSTT_").Append(ComponentId) : FString("STT_").Append(ComponentId);

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

	// execution_config
	{
		TSharedPtr<FJsonObject> ExecProps = MakeShared<FJsonObject>();
		ExecProps->SetStringField(TEXT("stt_component_id"), STTComponentId);
		TSharedPtr<FJsonObject> ExecCfg = MakeShared<FJsonObject>();
		ExecCfg->SetStringField(TEXT("type"), TEXT("STTNodeExecutionConfig"));
		ExecCfg->SetObjectField(TEXT("properties"), ExecProps);
		NodeObj->SetObjectField(TEXT("execution_config"), ExecCfg);
	}

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

	// creation_config
	{
		TArray<TSharedPtr<FJsonValue>> Components = GraphJson->GetArrayField(TEXT("components"));
		const TCHAR* CompType = ExecutionConfig.bStreamed ? TEXT("StreamingSTTInterface") : TEXT("STTInterface");
		bool bExists = false;
		for (const TSharedPtr<FJsonValue>& Component : Components)
		{
			const TSharedPtr<FJsonObject> ComponentObj = Component->AsObject();
			if (ComponentObj.IsValid() && ComponentObj->GetStringField(TEXT("id")) == STTComponentId)
			{
				bExists = true;
				break;
			}
		}
		if (!bExists)
		{
			TSharedPtr<FJsonObject> Comp = MakeShared<FJsonObject>();
			Comp->SetStringField(TEXT("id"), STTComponentId);
			Comp->SetStringField(TEXT("type"), CompType);
			// creation_config
			const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>();
			if (Settings)
			{
				TSharedPtr<FJsonObject> CreationCfg = MakeShared<FJsonObject>();
				TSharedPtr<FJsonObject> CreationProps = MakeShared<FJsonObject>();
				if (ExecutionConfig.bStreamed)
				{
					if (const FInworldStreamingSTTCreationConfig* Cfg = Settings->StreamingSTTCreationConfig.Find(ComponentId))
					{
						CreationCfg->SetStringField(TEXT("type"), Cfg->ComputeHost == EInworldComputeHost::REMOTE
																	  ? TEXT("RemoteStreamingSTTConfig")
																	  : TEXT("LocalStreamingSTTConfig"));
						if (Cfg->ComputeHost == EInworldComputeHost::REMOTE)
						{
							CreationProps->SetStringField(TEXT("api_key"), TEXT("{{INWORLD_API_KEY}}"));
						}
						if (Cfg->ComputeHost == EInworldComputeHost::LOCAL)
						{
							CreationProps->SetStringField(TEXT("stt_model_path"),
								FPaths::Combine("{{INWORLD_EXTRAS_PATH}}", "stt", Cfg->LocalCreationConfig.Model));
							CreationProps->SetStringField(TEXT("vad_model_path"),
								FPaths::Combine("{{INWORLD_EXTRAS_PATH}}", "vad", "silero", "silero_vad.onnx"));
							TSharedPtr<FJsonObject> Dev = MakeShared<FJsonObject>();
							Dev->SetStringField(TEXT("type"),
								Cfg->LocalCreationConfig.LocalComputeHost == EInworldLocalComputeHost::CPU ? TEXT("CPU")
																										   : TEXT("CUDA"));
							CreationProps->SetObjectField(TEXT("stt_device"), Dev);
							CreationProps->SetObjectField(TEXT("vad_device"), Dev);
						}
						TSharedPtr<FJsonObject> Def = MakeShared<FJsonObject>();
						if (Cfg->ComputeHost == EInworldComputeHost::LOCAL)
						{
							// TODO for 1.2.1 we should chnage to false
							bool bHasSpeechConfig = true;
							TSharedPtr<FJsonObject> Speech = MakeShared<FJsonObject>();
							if (Cfg->LocalCreationConfig.StreamSpeechRecognitionConfig.SpeechRecognitionConfig.bHasLanguage)
							{
								FString Code;
								inworld::Convert(
									Cfg->LocalCreationConfig.StreamSpeechRecognitionConfig.SpeechRecognitionConfig.Language, Code);
								Speech->SetStringField(TEXT("language_code"), Code);
								bHasSpeechConfig = true;
							}
							if (bHasSpeechConfig)
							{
								Def->SetObjectField(TEXT("speech_config"), Speech);
							}
							Def->SetNumberField(TEXT("silence_threshold_ms"),
								Cfg->LocalCreationConfig.StreamSpeechRecognitionConfig.SilenceThresholdMs);
						}
						else
						{
							// TODO for 1.2.1 we should chnage to false
							bool bHasSpeechConfig = true;
							TSharedPtr<FJsonObject> Speech = MakeShared<FJsonObject>();
							if (Cfg->RemoteCreationConfig.StreamSpeechRecognitionConfig.SpeechRecognitionConfig.bHasLanguage)
							{
								FString Code;
								inworld::Convert(
									Cfg->RemoteCreationConfig.StreamSpeechRecognitionConfig.SpeechRecognitionConfig.Language, Code);
								Speech->SetStringField(TEXT("language_code"), Code);
								bHasSpeechConfig = true;
							}
							if (bHasSpeechConfig)
							{
								Def->SetObjectField(TEXT("speech_config"), Speech);
							}
							Def->SetNumberField(TEXT("silence_threshold_ms"),
								Cfg->RemoteCreationConfig.StreamSpeechRecognitionConfig.SilenceThresholdMs);
						}
						CreationProps->SetObjectField(TEXT("default_config"), Def);
					}
				}
				else
				{
					if (const FInworldSTTCreationConfig* Cfg = Settings->STTCreationConfig.Find(ComponentId))
					{
						CreationCfg->SetStringField(TEXT("type"),
							Cfg->ComputeHost == EInworldComputeHost::REMOTE ? TEXT("RemoteSTTConfig") : TEXT("LocalSTTConfig"));
						if (Cfg->ComputeHost == EInworldComputeHost::REMOTE)
						{
							CreationProps->SetStringField(TEXT("api_key"), TEXT("{{INWORLD_API_KEY}}"));
						}
						if (Cfg->ComputeHost == EInworldComputeHost::LOCAL)
						{
							CreationProps->SetStringField(TEXT("model_path"),
								FPaths::Combine("{{INWORLD_EXTRAS_PATH}}", "embedders", Cfg->LocalCreationConfig.Model));
							TSharedPtr<FJsonObject> Dev = MakeShared<FJsonObject>();
							Dev->SetStringField(TEXT("type"),
								Cfg->LocalCreationConfig.LocalComputeHost == EInworldLocalComputeHost::CPU ? TEXT("CPU")
																										   : TEXT("CUDA"));
							CreationProps->SetObjectField(TEXT("device"), Dev);
						}
						TSharedPtr<FJsonObject> DefaultConfig = MakeShared<FJsonObject>();
						if (Cfg->ComputeHost == EInworldComputeHost::LOCAL)
						{
							if (Cfg->LocalCreationConfig.SpeechRecognitionConfig.bHasLanguage)
							{
								FString Code;
								inworld::Convert(Cfg->LocalCreationConfig.SpeechRecognitionConfig.Language, Code);
								DefaultConfig->SetStringField(TEXT("language_code"), Code);
							}
						}
						else
						{
							if (Cfg->RemoteCreationConfig.SpeechRecognitionConfig.bHasLanguage)
							{
								FString Code;
								inworld::Convert(Cfg->RemoteCreationConfig.SpeechRecognitionConfig.Language, Code);
								DefaultConfig->SetStringField(TEXT("language_code"), Code);
							}
						}
						CreationProps->SetObjectField(TEXT("default_config"), DefaultConfig);
					}
				}
				CreationCfg->SetObjectField(TEXT("properties"), CreationProps);
				Comp->SetObjectField(TEXT("creation_config"), CreationCfg);
			}
			Components.Add(MakeShared<FJsonValueObject>(Comp));
			GraphJson->SetArrayField(TEXT("components"), Components);
		}
	}
	return EInworldRuntimeErrors::SUCCESS;
}

UInworldNode_STT* UInworldNode_STT::CreateNative(
	class UObject* Outer, const FString& NodeName, const FInworldSTTNodeExecutionConfig& ExecutionConfig)
{
	UInworldNode_STT* NewNode = NewObject<UInworldNode_STT>(Outer, FName(NodeName));
	NewNode->NodeName = NodeName;
	NewNode->ExecutionConfig = ExecutionConfig;
	return NewNode;
}
