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

#include "Primitives/LLM/InworldLLMTasks.h"

#include "Interfaces/IPluginManager.h"
#include "InworldRuntimeUtils.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

static TVariant<TSharedPtr<inworld::LLMInterface>, FString> CreateLLM(const FInworldLocalLLMCreationConfig& LocalLLMCreationConfig)
{
#define RETURN_ERROR(Error) return TVariant<TSharedPtr<inworld::LLMInterface>, FString>{TInPlaceType<FString>(), Error};
#define RETURN_SUCCESS(Object) \
	return TVariant<TSharedPtr<inworld::LLMInterface>, FString>{TInPlaceType<TSharedPtr<inworld::LLMInterface>>(), Object};

	if (const FString ModelFilePath{GetInworldExtrasPath("llm", LocalLLMCreationConfig.Model)}; !FPaths::FileExists(ModelFilePath))
	{
		RETURN_ERROR(FString::Format(TEXT("{0}: Model not found at {1}"), {*FString(__func__), *ModelFilePath}))
	}

	inworld::LocalLLMConfig local_llm_config;
	inworld::Convert(LocalLLMCreationConfig, local_llm_config);
	if (const inworld::StatusOr_LLMInterface status_or_llm{inworld::LLMFactory().CreateLLM(local_llm_config)}; status_or_llm.ok())
	{
		RETURN_SUCCESS(MakeShared<inworld::LLMInterface>(status_or_llm.value()));
	}
	else
	{
		RETURN_ERROR(
			FString::Format(TEXT("{0}: {1}"), {*FString(__func__), UTF8_TO_TCHAR(status_or_llm.status().ToString().c_str())}))
	}
#undef RETURN_ERROR
#undef RETURN_SUCCESS
}

static TVariant<TSharedPtr<inworld::LLMInterface>, FString> CreateLLM(
	const FInworldRemoteLLMCreationConfig& RemoteLLMCreationConfig)
{
#define RETURN_ERROR(Error) return TVariant<TSharedPtr<inworld::LLMInterface>, FString>{TInPlaceType<FString>(), Error};
#define RETURN_SUCCESS(Object) \
	return TVariant<TSharedPtr<inworld::LLMInterface>, FString>{TInPlaceType<TSharedPtr<inworld::LLMInterface>>(), Object};

	const FString& RuntimeApiKey{GetInworldRuntimeApiKey()};
	if (RuntimeApiKey.IsEmpty())
	{
		RETURN_ERROR(FString::Format(TEXT("{0}: Inworld Runtime API Key is not set. Check ProjectSettings."), {*FString(__func__)}))
	}

	inworld::RemoteLLMConfig remote_llm_config;
	inworld::Convert(RemoteLLMCreationConfig, remote_llm_config);
	if (const inworld::StatusOr_LLMInterface status_or_llm{inworld::LLMFactory().CreateLLM(remote_llm_config)}; status_or_llm.ok())
	{
		RETURN_SUCCESS(MakeShared<inworld::LLMInterface>(status_or_llm.value()));
	}
	else
	{
		RETURN_ERROR(
			FString::Format(TEXT("{0}: {1}"), {*FString(__func__), UTF8_TO_TCHAR(status_or_llm.status().ToString().c_str())}))
	}
#undef RETURN_ERROR
#undef RETURN_SUCCESS
}

namespace inworld
{
namespace Tasks
{
namespace LLM
{

UE::Tasks::TTask<TVariant<TSharedPtr<inworld::LLMInterface>, FString>> CreateLLMTask(
	const FInworldLLMCreationConfig& LLMCreationConfig)
{
#define RETURN_ERROR(Error) return TVariant<TSharedPtr<inworld::LLMInterface>, FString>{TInPlaceType<FString>(), Error};

	return UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[LLMCreationConfig]()
		{
			switch (LLMCreationConfig.ComputeHost)
			{
				case EInworldComputeHost::LOCAL:
				{
					const FInworldLocalLLMCreationConfig& LocalLLMCreationConfig{LLMCreationConfig.LocalCreationConfig};
					return CreateLLM(LocalLLMCreationConfig);
				}
				break;
				case EInworldComputeHost::REMOTE:
				{
					const FInworldRemoteLLMCreationConfig& RemoteLLMCreationConfig{LLMCreationConfig.RemoteCreationConfig};
					return CreateLLM(RemoteLLMCreationConfig);
				}
				break;
				default:
					RETURN_ERROR(
						FString::Format(TEXT("{0}: Invalid Enum Value: LLMCreationConfig.ComputeHost"), {*FString(__func__)}));
			}
		});
#undef RETURN_ERROR
}

}	 // namespace LLM
}	 // namespace Tasks
}	 // namespace inworld
