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

#include "Primitives/VAD/InworldVADTasks.h"

#include "Graph/InworldData.h"
#include "InworldRuntimeTypes.h"
#include "InworldRuntimeUtils.h"
#include "Misc/Paths.h"
#include "Modules/ModuleManager.h"
#include "Primitives/VAD/InworldVADTypes.h"

#include <Interfaces/IPluginManager.h>

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

static TVariant<TSharedPtr<inworld::VADInterface>, FString> CreateVAD()
{
#define RETURN_ERROR(Error) return TVariant<TSharedPtr<inworld::VADInterface>, FString>{TInPlaceType<FString>(), Error};
#define RETURN_SUCCESS(Object) \
	return TVariant<TSharedPtr<inworld::VADInterface>, FString>{TInPlaceType<TSharedPtr<inworld::VADInterface>>(), Object};

#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION < 6
	if (FModuleManager::Get().IsModuleLoaded("NNEDenoiser") || FModuleManager::Get().IsModuleLoaded("NNERuntimeORT"))
	{
		RETURN_ERROR(
			FString::Format(TEXT("{0}: VAD is not compatible with NNEDenoiser and NNERuntimeORT Plugins."), {*FString(__func__)}))
	}
#endif

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

	inworld::LocalVADConfig local_vad_config;
	inworld::Convert(FInworldVADCreationConfig{}, local_vad_config);

	if (const inworld::StatusOr_VADInterface status_or_vad{inworld::VADFactory().CreateVAD(local_vad_config)}; status_or_vad.ok())
	{
		RETURN_SUCCESS(MakeShared<inworld::VADInterface>(status_or_vad.value()))
	}
	else
	{
		RETURN_ERROR(
			FString::Format(TEXT("{0}: {1}"), {*FString(__func__), UTF8_TO_TCHAR(status_or_vad.status().ToString().c_str())}))
	}
#undef RETURN_ERROR
#undef RETURN_SUCCESS
}

namespace inworld
{
namespace Tasks
{
namespace VAD
{

UE::Tasks::TTask<TVariant<TSharedPtr<inworld::VADInterface>, FString>> CreateVADTask(const FInworldVADCreationConfig& VADConfig)
{
	return UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[]()
		{
			return CreateVAD();
		});
}

UE::Tasks::TTask<TVariant<bool, FString>> CreateDetectVoiceActivityTask(
	TSharedPtr<inworld::VADInterface> VADInterface, const FInworldVADConfig& Config, const FInworldData_Audio& Audio)
{
#define RETURN_ERROR(Error) return TVariant<bool, FString>{TInPlaceType<FString>(), Error};
#define RETURN_SUCCESS(Object) return TVariant<bool, FString>{TInPlaceType<bool>(), Object};

	return UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[VADInterface, Config, Audio]()
		{
			if (!VADInterface.IsValid() || !VADInterface->is_valid())
			{
				RETURN_ERROR(TEXT("Invalid VAD interface"));
			}

			inworld::VoiceActivityDetectionConfig voice_activity_detection_config;
			voice_activity_detection_config.speech_threshold(Config.SpeechThreshold);

			inworld::AudioChunk audio_chunk;
			audio_chunk.sample_rate(Audio.SampleRate);
			inworld::vector_float vector_float =
				inworld::memcpy_to_vector_float((float*) Audio.Waveform->GetData(), Audio.Waveform->Num() / sizeof(float));
			audio_chunk.data(&vector_float);

			if (const inworld::StatusOr_int status_or_vad_result{
					VADInterface->DetectVoiceActivity(audio_chunk, voice_activity_detection_config)};
				status_or_vad_result.ok())
			{
				const bool bVoiceActivityDetected{status_or_vad_result.value() != INDEX_NONE};
				RETURN_SUCCESS(bVoiceActivityDetected)
			}
			else
			{
				RETURN_ERROR(FString::Format(
					TEXT("{0}: {1}"), {*FString(__func__), UTF8_TO_TCHAR(status_or_vad_result.status().ToString().c_str())}))
			}
		});

#undef RETURN_ERROR
#undef RETURN_SUCCESS
}

}	 // namespace VAD
}	 // namespace Tasks
}	 // namespace inworld
