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

#include "Graph/Nodes/InworldNode_TTS.h"

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

UInworldNode_TTS::UInworldNode_TTS()
{
	NodeName = "TTS";
}

UInworldNode_TTS* UInworldNode_TTS::CreateNative(
	UObject* Outer, const FString& NodeName, const FInworldTTSNodeExecutionConfig& ExecutionConfig)
{
	UInworldNode_TTS* Node = NewObject<UInworldNode_TTS>(Outer, FName(NodeName));
	Node->NodeName = NodeName;
	Node->ExecutionConfig = ExecutionConfig;
	return Node;
}

EInworldRuntimeErrors UInworldNode_TTS::GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson)
{
	const FString TTSComponentId = FString("TTS_").Append(ExecutionConfig.TTSComponentId);

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

	// execution_config
	{
		TSharedPtr<FJsonObject> ExecProps = MakeShared<FJsonObject>();
		ExecProps->SetStringField(TEXT("tts_component_id"), TTSComponentId);

		// voice
		{
			TSharedPtr<FJsonObject> Voice = MakeShared<FJsonObject>();
			Voice->SetStringField(TEXT("id"), ExecutionConfig.Voice.Id);
			if (ExecutionConfig.Voice.bHasLanguage)
			{
				FString LangCode = TEXT("");
				inworld::Convert(ExecutionConfig.Voice.Language, LangCode);
				if (!LangCode.IsEmpty())
				{
					Voice->SetStringField(TEXT("language_code"), LangCode);
				}
			}
			ExecProps->SetObjectField(TEXT("voice"), Voice);
		}

		// synthesis_config
		if (ExecutionConfig.bOverrideSpeechSynthesisConfig)
		{
			TSharedPtr<FJsonObject> Syn = MakeShared<FJsonObject>();
			if (ExecutionConfig.Provider == EInworldRemoteTTSProvider::INWORLD)
			{
				TSharedPtr<FJsonObject> SynCfg = MakeShared<FJsonObject>();
				SynCfg->SetStringField(TEXT("model_id"), ExecutionConfig.InworldSpeechSynthesisConfig.ModelId);
				// postprocessing
				{
					TSharedPtr<FJsonObject> Post = MakeShared<FJsonObject>();
					Post->SetNumberField(TEXT("sample_rate"),
						ExecutionConfig.InworldSpeechSynthesisConfig.SpeechSynthesisPostprocessingConfig.SampleRate);
					SynCfg->SetObjectField(TEXT("postprocessing"), Post);
				}
				// inference
				{
					TSharedPtr<FJsonObject> Inf = MakeShared<FJsonObject>();
					if (ExecutionConfig.InworldSpeechSynthesisConfig.SpeechSynthesisInferenceConfig.bOverrideTemperature)
					{
						Inf->SetNumberField(TEXT("temperature"),
							ExecutionConfig.InworldSpeechSynthesisConfig.SpeechSynthesisInferenceConfig.Temperature);
					}
					Inf->SetNumberField(TEXT("speaking_rate"),
						ExecutionConfig.InworldSpeechSynthesisConfig.SpeechSynthesisInferenceConfig.SpeakingRate);
					SynCfg->SetObjectField(TEXT("inference"), Inf);
				}
				Syn->SetStringField(TEXT("type"), TEXT("inworld"));
				Syn->SetObjectField(TEXT("config"), SynCfg);
			}
			else if (ExecutionConfig.Provider == EInworldRemoteTTSProvider::ELEVEN_LABS)
			{
				TSharedPtr<FJsonObject> SynCfg = MakeShared<FJsonObject>();
				// query_params
				{
					TSharedPtr<FJsonObject> Query = MakeShared<FJsonObject>();
					Query->SetBoolField(TEXT("enable_logging"), ExecutionConfig.ElevenLabsSpeechSynthesisConfig.EnableLogging);
					// output_format enum to string
					FString OutputFmt{};
					inworld::Convert(ExecutionConfig.ElevenLabsSpeechSynthesisConfig.OutputFormat, OutputFmt);
					Query->SetStringField(TEXT("output_format"), OutputFmt);
					SynCfg->SetObjectField(TEXT("query_params"), Query);
				}
				// request_params
				{
					TSharedPtr<FJsonObject> Req = MakeShared<FJsonObject>();
					Req->SetStringField(TEXT("model_id"), ExecutionConfig.ElevenLabsSpeechSynthesisConfig.ModelId);
					FString Norm = TEXT("AUTO");
					switch (ExecutionConfig.ElevenLabsSpeechSynthesisConfig.TextNormalization)
					{
						case EElevenLabsTTSTextNormalization::AUTO:
							Norm = TEXT("AUTO");
							break;
						case EElevenLabsTTSTextNormalization::ON:
							Norm = TEXT("ON");
							break;
						case EElevenLabsTTSTextNormalization::OFF:
							Norm = TEXT("OFF");
							break;
						default:
							break;
					}
					Req->SetStringField(TEXT("apply_text_normalization"), Norm);
					SynCfg->SetObjectField(TEXT("request_params"), Req);
				}
				Syn->SetStringField(TEXT("type"), TEXT("elevenlabs"));
				Syn->SetObjectField(TEXT("config"), SynCfg);
			}
			ExecProps->SetObjectField(TEXT("synthesis_config"), Syn);
		}

		TSharedPtr<FJsonObject> ExecCfg = MakeShared<FJsonObject>();
		ExecCfg->SetStringField(TEXT("type"), TEXT("TTSNodeExecutionConfig"));
		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);

	{
		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")) == TTSComponentId)
			{
				bExists = true;
				break;
			}
		}
		if (!bExists)
		{
			TSharedPtr<FJsonObject> Component = MakeShared<FJsonObject>();
			Component->SetStringField(TEXT("id"), TTSComponentId);
			Component->SetStringField(TEXT("type"), TEXT("TTSInterface"));
			// creation_config
			const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>();
			if (Settings)
			{
				if (const FInworldTTSCreationConfig* TtsCfg = Settings->TTSCreationConfig.Find(ExecutionConfig.TTSComponentId))
				{
					TSharedPtr<FJsonObject> CreationCfg = MakeShared<FJsonObject>();
					TSharedPtr<FJsonObject> CreationProps = MakeShared<FJsonObject>();
					if (TtsCfg->ComputeHost == EInworldComputeHost::REMOTE)
					{
						CreationCfg->SetStringField(TEXT("type"), TEXT("RemoteTTSConfig"));
						CreationProps->SetStringField(TEXT("api_key"), TEXT("{{INWORLD_API_KEY}}"));
						TSharedPtr<FJsonObject> Syn = MakeShared<FJsonObject>();
						if (TtsCfg->RemoteCreationConfig.Provider == EInworldRemoteTTSProvider::INWORLD)
						{
							TSharedPtr<FJsonObject> SynCfg = MakeShared<FJsonObject>();
							SynCfg->SetStringField(
								TEXT("model_id"), TtsCfg->RemoteCreationConfig.InworldSpeechSynthesisConfig.ModelId);
							// postprocessing
							{
								TSharedPtr<FJsonObject> Post = MakeShared<FJsonObject>();
								Post->SetNumberField(TEXT("sample_rate"), TtsCfg->RemoteCreationConfig.InworldSpeechSynthesisConfig
																			  .SpeechSynthesisPostprocessingConfig.SampleRate);
								SynCfg->SetObjectField(TEXT("postprocessing"), Post);
							}
							// inference
							{
								TSharedPtr<FJsonObject> Inf = MakeShared<FJsonObject>();
								if (TtsCfg->RemoteCreationConfig.InworldSpeechSynthesisConfig.SpeechSynthesisInferenceConfig
										.bOverrideTemperature)
								{
									Inf->SetNumberField(
										TEXT("temperature"), TtsCfg->RemoteCreationConfig.InworldSpeechSynthesisConfig
																 .SpeechSynthesisInferenceConfig.Temperature);
								}
								Inf->SetNumberField(TEXT("speaking_rate"), TtsCfg->RemoteCreationConfig.InworldSpeechSynthesisConfig
																			   .SpeechSynthesisInferenceConfig.SpeakingRate);
								SynCfg->SetObjectField(TEXT("inference"), Inf);
							}
							Syn->SetStringField(TEXT("type"), TEXT("inworld"));
							Syn->SetObjectField(TEXT("config"), SynCfg);
						}
						else if (TtsCfg->RemoteCreationConfig.Provider == EInworldRemoteTTSProvider::ELEVEN_LABS)
						{
							TSharedPtr<FJsonObject> SynCfg = MakeShared<FJsonObject>();
							// query_params
							{
								TSharedPtr<FJsonObject> Query = MakeShared<FJsonObject>();
								Query->SetBoolField(TEXT("enable_logging"),
									TtsCfg->RemoteCreationConfig.ElevenLabsSpeechSynthesisConfig.EnableLogging);
								// output_format enum to string
								FString OutputFmt{};
								inworld::Convert(
									TtsCfg->RemoteCreationConfig.ElevenLabsSpeechSynthesisConfig.OutputFormat, OutputFmt);
								Query->SetStringField(TEXT("output_format"), OutputFmt);
								SynCfg->SetObjectField(TEXT("query_params"), Query);
							}
							// request_params
							{
								TSharedPtr<FJsonObject> Req = MakeShared<FJsonObject>();
								Req->SetStringField(
									TEXT("model_id"), TtsCfg->RemoteCreationConfig.ElevenLabsSpeechSynthesisConfig.ModelId);
								FString Norm = TEXT("AUTO");
								switch (TtsCfg->RemoteCreationConfig.ElevenLabsSpeechSynthesisConfig.TextNormalization)
								{
									case EElevenLabsTTSTextNormalization::AUTO:
										Norm = TEXT("AUTO");
										break;
									case EElevenLabsTTSTextNormalization::ON:
										Norm = TEXT("ON");
										break;
									case EElevenLabsTTSTextNormalization::OFF:
										Norm = TEXT("OFF");
										break;
									default:
										break;
								}
								Req->SetStringField(TEXT("apply_text_normalization"), Norm);
								SynCfg->SetObjectField(TEXT("request_params"), Req);
							}
							Syn->SetStringField(TEXT("type"), TEXT("elevenlabs"));
							Syn->SetObjectField(TEXT("config"), SynCfg);
						}
						CreationProps->SetObjectField(TEXT("synthesis_config"), Syn);
					}
					else
					{
						return EInworldRuntimeErrors::LOCAL_TTS_NOTSUPPORTED;
						// Local TTS not supported
					}
					CreationCfg->SetObjectField(TEXT("properties"), CreationProps);
					Component->SetObjectField(TEXT("creation_config"), CreationCfg);
				}
			}
			Components.Add(MakeShared<FJsonValueObject>(Component));
			GraphJson->SetArrayField(TEXT("components"), Components);
		}
	}
	return EInworldRuntimeErrors::SUCCESS;
}

UInworldNode_TTSRequest::UInworldNode_TTSRequest()
{
	NodeName = "TTS Request";
}

FInworldData_TTSRequest UInworldNode_TTSRequest::Process_DataStream_String(const UInworldProcessContext* ProcessContext,
	const FInworldData_DataStream_String& TextStream, const FInworldVoice& Voice) const
{
	FInworldData_TTSRequest TTSRequest;
	TTSRequest.bHasStreamContent = true;
	TTSRequest.InputStream = TextStream.InputStream;
	TTSRequest.Voice = Voice;
	return TTSRequest;
}

FInworldData_TTSRequest UInworldNode_TTSRequest::Process_Text(
	const UInworldProcessContext* ProcessContext, const FInworldData_Text& Text, const FInworldVoice& Voice) const
{
	FInworldData_TTSRequest TTSRequest;
	TTSRequest.bHasTextContent = true;
	TTSRequest.Text = Text.Text;
	TTSRequest.Voice = Voice;
	return TTSRequest;
}
