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

#include "Graph/Nodes/InworldNode_LLM.h"

#include "Dom/JsonObject.h"
#include "Graph/InworldData.h"
#include "Graph/InworldGraph.h"
#include "InworldBlueprintFunctionLibrary.h"
#include "InworldRuntimeSettings.h"
#include "InworldRuntimeSubsystem.h"
#include "JsonObjectConverter.h"
#include "Policies/PrettyJsonPrintPolicy.h"
#include "Serialization/JsonSerializer.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

UInworldNode_LLM::UInworldNode_LLM()
{
	NodeName = "LLM";
}

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

EInworldRuntimeErrors UInworldNode_LLM::GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson)
{
	TSharedPtr<FJsonObject> NodeObj = MakeShared<FJsonObject>();
	NodeObj->SetStringField(TEXT("id"), GetNodeName());
	NodeObj->SetStringField(TEXT("type"), TEXT("LLMChatNode"));

	// execution_config
	{
		TSharedPtr<FJsonObject> ExecProps = MakeShared<FJsonObject>();
		ExecProps->SetStringField(TEXT("llm_component_id"), ExecutionConfig.LLMComponentId);

		// text_generation_config
		{
			TSharedPtr<FJsonObject> Gen = MakeShared<FJsonObject>();
			Gen->SetNumberField(TEXT("max_new_tokens"), ExecutionConfig.TextGenerationConfig.MaxNewTokens);
			Gen->SetNumberField(TEXT("max_prompt_length"), ExecutionConfig.TextGenerationConfig.MaxPromptLength);
			Gen->SetNumberField(TEXT("temperature"), ExecutionConfig.TextGenerationConfig.Temperature);
			Gen->SetNumberField(TEXT("top_p"), ExecutionConfig.TextGenerationConfig.TopP);
			Gen->SetNumberField(TEXT("repetition_penalty"), ExecutionConfig.TextGenerationConfig.RepetitionPenalty);
			Gen->SetNumberField(TEXT("frequency_penalty"), ExecutionConfig.TextGenerationConfig.FrequencyPenalty);
			Gen->SetNumberField(TEXT("presence_penalty"), ExecutionConfig.TextGenerationConfig.PresencePenalty);
			TArray<TSharedPtr<FJsonValue>> Stops;
			for (const FString& S : ExecutionConfig.TextGenerationConfig.StopSequences)
			{
				Stops.Add(MakeShared<FJsonValueString>(S));
			}
			Gen->SetArrayField(TEXT("stop_sequences"), Stops);
			ExecProps->SetObjectField(TEXT("text_generation_config"), Gen);
		}

		// message_templates
		{
			if (ExecutionConfig.MessageTemplates.Num() > 0)
			{
				TArray<TSharedPtr<FJsonValue>> MsgTemplates;
				for (const FInworldLLMMessage& Msg : ExecutionConfig.MessageTemplates)
				{
					TSharedPtr<FJsonObject> MsgTemplate = MakeShared<FJsonObject>();
					switch (Msg.Role)
					{
						case EInworldLLMMessageRole::USER:
							MsgTemplate->SetStringField(TEXT("role"), TEXT("user"));
							break;
						case EInworldLLMMessageRole::ASSISTANT:
							MsgTemplate->SetStringField(TEXT("role"), TEXT("assistant"));
							break;
						case EInworldLLMMessageRole::TOOL:
							MsgTemplate->SetStringField(TEXT("role"), TEXT("tool"));
							break;
						case EInworldLLMMessageRole::SYSTEM:
							MsgTemplate->SetStringField(TEXT("role"), TEXT("system"));
							break;
						case EInworldLLMMessageRole::DEVELOPER:
							MsgTemplate->SetStringField(TEXT("role"), TEXT("developer"));
							break;
					}
					TSharedPtr<FJsonObject> Content = MakeShared<FJsonObject>();
					switch (Msg.ContentType)
					{
						case EInworldChatMessageContentType::TEXT:
							Content->SetStringField(TEXT("type"), TEXT("text"));
							Content->SetStringField(TEXT("value"), Msg.Content);
							break;
						case EInworldChatMessageContentType::TEMPLATE:
							Content->SetStringField(TEXT("type"), TEXT("template"));
							Content->SetStringField(TEXT("template"), Msg.Content);
							break;
					}
					MsgTemplate->SetObjectField(TEXT("content"), Content);
					MsgTemplates.Add(MakeShareable(new FJsonValueObject(MsgTemplate)));
				}

				ExecProps->SetArrayField(TEXT("message_templates"), MsgTemplates);
			}
		}

		ExecProps->SetBoolField(TEXT("stream"), ExecutionConfig.bStream);

		TSharedPtr<FJsonObject> ExecCfg = MakeShared<FJsonObject>();
		ExecCfg->SetStringField(TEXT("type"), TEXT("LLMChatNodeExecutionConfig"));
		ExecCfg->SetObjectField(TEXT("properties"), ExecProps);
		NodeObj->SetObjectField(TEXT("execution_config"), ExecCfg);
	}

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

	// LLM component
	{
		TArray<TSharedPtr<FJsonValue>> Components = GraphJson->GetArrayField(TEXT("components"));
		const FString ComponentId = ExecutionConfig.LLMComponentId;
		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("LLMInterface"));

			// component + creation_config
			const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>();
			if (Settings)
			{
				if (const FInworldLLMCreationConfig* LlmCfg = Settings->LLMCreationConfig.Find(ComponentId))
				{
					TSharedPtr<FJsonObject> CreationCfg = MakeShared<FJsonObject>();
					TSharedPtr<FJsonObject> CreationProps = MakeShared<FJsonObject>();
					if (LlmCfg->ComputeHost == EInworldComputeHost::REMOTE)
					{
						CreationCfg->SetStringField(TEXT("type"), TEXT("RemoteLLMConfig"));
						CreationProps->SetStringField(TEXT("provider"), LlmCfg->RemoteCreationConfig.Provider);
						CreationProps->SetStringField(TEXT("model_name"), LlmCfg->RemoteCreationConfig.Model);
						CreationProps->SetStringField(TEXT("api_key"), TEXT("{{INWORLD_API_KEY}}"));
					}
					else
					{
						CreationCfg->SetStringField(TEXT("type"), TEXT("LocalLLMConfig"));
						CreationProps->SetStringField(TEXT("model_path"),
							FPaths::Combine("{{INWORLD_EXTRAS_PATH}}", "embedders", LlmCfg->LocalCreationConfig.Model));

						TSharedPtr<FJsonObject> DeviceObj = MakeShared<FJsonObject>();
						const TCHAR* DeviceType = (LlmCfg->LocalCreationConfig.LocalComputeHost == EInworldLocalComputeHost::CPU)
													  ? TEXT("CPU")
													  : TEXT("CUDA");
						DeviceObj->SetStringField(TEXT("type"), DeviceType);
						CreationProps->SetObjectField(TEXT("device"), DeviceObj);
					}

					// ExecutionConfig.TextGenerationConfig
					TSharedPtr<FJsonObject> DefaultCfg = MakeShared<FJsonObject>();
					DefaultCfg->SetNumberField(TEXT("max_new_tokens"), ExecutionConfig.TextGenerationConfig.MaxNewTokens);
					DefaultCfg->SetNumberField(TEXT("max_prompt_length"), ExecutionConfig.TextGenerationConfig.MaxPromptLength);
					DefaultCfg->SetNumberField(TEXT("temperature"), ExecutionConfig.TextGenerationConfig.Temperature);
					DefaultCfg->SetNumberField(TEXT("top_p"), ExecutionConfig.TextGenerationConfig.TopP);
					DefaultCfg->SetNumberField(TEXT("repetition_penalty"), ExecutionConfig.TextGenerationConfig.RepetitionPenalty);
					DefaultCfg->SetNumberField(TEXT("frequency_penalty"), ExecutionConfig.TextGenerationConfig.FrequencyPenalty);
					DefaultCfg->SetNumberField(TEXT("presence_penalty"), ExecutionConfig.TextGenerationConfig.PresencePenalty);
					TArray<TSharedPtr<FJsonValue>> StopSequences;
					for (const FString& StopSequence : ExecutionConfig.TextGenerationConfig.StopSequences)
					{
						StopSequences.Add(MakeShared<FJsonValueString>(StopSequence));
					}
					DefaultCfg->SetArrayField(TEXT("stop_sequences"), StopSequences);
					CreationProps->SetObjectField(TEXT("default_config"), DefaultCfg);

					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_LLMToTextStream::UInworldNode_LLMToTextStream()
{
	NodeName = "LLMResponse To TextStream";

	bExecuteProcessInGameThread = false;
}

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

FInworldData_DataStream_String UInworldNode_LLMToTextStream::Process_Default(const FInworldData_LLMChatResponse& ChatResponse) const
{
	if (TSharedPtr<inworld::InputStream_Content> Stream = ChatResponse.InputStream; Stream.IsValid())
	{
		if (inworld::StatusOr_ContentToStringStream status_or_content_to_string_stream{
				inworld::ContentToStringStream::Create(*(Stream.Get()))};
			status_or_content_to_string_stream.ok())
		{
			FInworldData_DataStream_String TextDataStream;
			TextDataStream.InputStream = MakeShared<inworld::InputStream_string>(status_or_content_to_string_stream.value());
			return TextDataStream;
		}
	}

	return FInworldData_DataStream_String::Error(TEXT("Invalid Input Stream."));
}

UInworldNode_LLMToText::UInworldNode_LLMToText()
{
	NodeName = "LLMResponse To Text";

	bExecuteProcessInGameThread = false;
}

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

FInworldData_Text UInworldNode_LLMToText::Process_Default(const FInworldData_LLMChatResponse& ChatResponse) const
{
	FInworldData_Text TextData;
	TextData.Text = ChatResponse.Content;
	return TextData;
}
