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

#include "Graph/Nodes/InworldNode_MCP.h"

#include "InworldRuntimeSettings.h"

static bool DoesMCPComponentExist(const TSharedRef<FJsonObject>& GraphJson, const FString& ComponentId)
{
	const FString MCPComponentId = FString("MCP_").Append(ComponentId);
	const TArray<TSharedPtr<FJsonValue>> Components = GraphJson->GetArrayField(TEXT("components"));
	for (const TSharedPtr<FJsonValue>& Component : Components)
	{
		const TSharedPtr<FJsonObject> ComponentObj = Component->AsObject();
		if (ComponentObj.IsValid() && ComponentObj->GetStringField(TEXT("id")) == MCPComponentId)
		{
			return true;
		}
	}
	return false;
}

static void CreateMCPComponent(const TSharedRef<FJsonObject>& GraphJson, const FString& ComponentId)
{
	const FString MCPComponentId = FString("MCP_").Append(ComponentId);
	TSharedPtr<FJsonObject> Comp = MakeShared<FJsonObject>();
	Comp->SetStringField(TEXT("id"), MCPComponentId);
	Comp->SetStringField(TEXT("type"), TEXT("MCPClientInterface"));

	if (const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>())
	{
		TSharedPtr<FJsonObject> CreationCfg = MakeShared<FJsonObject>();
		TSharedPtr<FJsonObject> CreationProps = MakeShared<FJsonObject>();

		if (const FInworldMCPClientCreationConfig* Cfg = Settings->MCPClientCreationConfig.Find(ComponentId))
		{
			CreationCfg->SetStringField(TEXT("type"), TEXT("MCPClientCreationConfig"));

			// SessionConfig
			{
				TSharedPtr<FJsonObject> SessionCfgObj = MakeShared<FJsonObject>();

				// Basic session fields
				const FInworldMCPSessionCreationConfig& SessionCfg = Cfg->SessionConfig;
				SessionCfgObj->SetStringField(TEXT("transport"), SessionCfg.Transport);
				SessionCfgObj->SetStringField(TEXT("endpoint"), SessionCfg.Endpoint);

				TSharedPtr<FJsonObject> TransportCfgObj = MakeShared<FJsonObject>();

				// Transport config
				switch (SessionCfg.TransportType)
				{
					// HTTP transport config
					case EInworldMCPTransport::Http:
					{
						TSharedPtr<FJsonObject> HttpCfgObj = MakeShared<FJsonObject>();
						HttpCfgObj->SetStringField(TEXT("api_key"), SessionCfg.HttpTransportConfig.APIKey);

						TransportCfgObj->SetStringField(TEXT("type"), TEXT("http"));
						TransportCfgObj->SetObjectField(TEXT("config"), HttpCfgObj);
					}
					break;

					// Studio transport config
					case EInworldMCPTransport::Studio:
					{
						TSharedPtr<FJsonObject> StudioCfgObj = MakeShared<FJsonObject>();

						// Environment as an object of key/value pairs
						TSharedPtr<FJsonObject> EnvObj = MakeShared<FJsonObject>();
						for (const TPair<FString, FString>& Pair : SessionCfg.StudioTransportConfig.Environment)
						{
							EnvObj->SetStringField(Pair.Key, Pair.Value);
						}
						StudioCfgObj->SetObjectField(TEXT("env"), EnvObj);

						TransportCfgObj->SetStringField(TEXT("type"), TEXT("stdio"));
						TransportCfgObj->SetObjectField(TEXT("config"), StudioCfgObj);
					}
					break;

					// Paragon transport config
					case EInworldMCPTransport::Paragon:
					{
						TSharedPtr<FJsonObject> ParagonCfgObj = MakeShared<FJsonObject>();
						ParagonCfgObj->SetStringField(TEXT("jwt_token"), SessionCfg.ParagonTransportConfig.JWTToken);

						TransportCfgObj->SetStringField(TEXT("type"), TEXT("paragon"));
						TransportCfgObj->SetObjectField(TEXT("config"), ParagonCfgObj);
					}
					break;
					default:
						break;
				}

				SessionCfgObj->SetObjectField(TEXT("auth_config"), TransportCfgObj);
				CreationProps->SetObjectField(TEXT("session_config"), SessionCfgObj);
			}
		}

		CreationCfg->SetObjectField(TEXT("properties"), CreationProps);
		Comp->SetObjectField(TEXT("creation_config"), CreationCfg);
	}

	TArray<TSharedPtr<FJsonValue>> Components = GraphJson->GetArrayField(TEXT("components"));
	Components.Add(MakeShared<FJsonValueObject>(Comp));
	GraphJson->SetArrayField(TEXT("components"), Components);
}

UInworldNode_MCPListTools::UInworldNode_MCPListTools()
{
	NodeName = "MCPListTools";
	bIsGetter = true;
}

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

EInworldRuntimeErrors UInworldNode_MCPListTools::GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson)
{
	const FString MCPComponentId = FString("MCP_").Append(ExecutionConfig.MCPComponentID);

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

	// execution_config
	{
		TSharedPtr<FJsonObject> ExecProps = MakeShared<FJsonObject>();
		ExecProps->SetStringField(TEXT("mcp_component_id"), MCPComponentId);
		TSharedPtr<FJsonObject> ExecCfg = MakeShared<FJsonObject>();
		ExecCfg->SetStringField(TEXT("type"), TEXT("MCPListToolsNodeExecutionConfig"));
		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
	if (!DoesMCPComponentExist(GraphJson, ExecutionConfig.MCPComponentID))
	{
		CreateMCPComponent(GraphJson, ExecutionConfig.MCPComponentID);
	}
	return EInworldRuntimeErrors::SUCCESS;
}

UInworldNode_MCPCallTool::UInworldNode_MCPCallTool()
{
	NodeName = "MCPCallTool";
}

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

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

	// execution_config
	{
		TSharedPtr<FJsonObject> ExecProps = MakeShared<FJsonObject>();
		ExecProps->SetStringField(TEXT("mcp_component_id"), ExecutionConfig.MCPComponentID);
		TSharedPtr<FJsonObject> ExecCfg = MakeShared<FJsonObject>();
		ExecCfg->SetStringField(TEXT("type"), TEXT("MCPCallToolNodeExecutionConfig"));
		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
	if (!DoesMCPComponentExist(GraphJson, ExecutionConfig.MCPComponentID))
	{
		CreateMCPComponent(GraphJson, ExecutionConfig.MCPComponentID);
	}
	return EInworldRuntimeErrors::SUCCESS;
}

UInworldNode_ExtractMCPRequest::UInworldNode_ExtractMCPRequest()
{
	NodeName = "ExtractMCPRequest";
}

FInworldData_ToolCallRequest UInworldNode_ExtractMCPRequest::Process_Default(const FInworldData_Text& Text) const
{
	FInworldData_ToolCallRequest Request;

	TSharedPtr<FJsonObject> JsonObject;
	const TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create(Text.Text);

	if (FJsonSerializer::Deserialize(Reader, JsonObject) && JsonObject.IsValid())
	{
		if (const TArray<TSharedPtr<FJsonValue>>* ToolCallsArray = nullptr;
			JsonObject->TryGetArrayField(TEXT("ToolCalls"), ToolCallsArray))
		{
			for (const TSharedPtr<FJsonValue>& Val : *ToolCallsArray)
			{
				if (const TSharedPtr<FJsonObject> ToolCallObj = Val->AsObject(); ToolCallObj.IsValid())
				{
					FInworldToolCall& ToolCall = Request.ToolCalls.AddDefaulted_GetRef();

					ToolCallObj->TryGetStringField(TEXT("id"), ToolCall.Id);
					ToolCallObj->TryGetStringField(TEXT("name"), ToolCall.Name);

					// Args is a JSON string, so we actually don't want this particular parameter deserialized further
					if (const TSharedPtr<FJsonObject>* ArgsObj = nullptr; ToolCallObj->TryGetObjectField(TEXT("args"), ArgsObj))
					{
						const TSharedRef<TJsonWriter<>> Writer = TJsonWriterFactory<>::Create(&ToolCall.Args);
						FJsonSerializer::Serialize(ArgsObj->ToSharedRef(), Writer);
					}
				}
			}

			return Request;
		}
	}

	Request.Error(FString::Format(TEXT("Failed to parse JSON: {0}"), {Text.Text}));
	return Request;
}
