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

#include "Primitives/AEC/InworldAECTasks.h"

#include "Graph/InworldData.h"
#include "InworldRuntimeTypes.h"
#include "InworldRuntimeUtils.h"
#include "Misc/Paths.h"
#include "Modules/ModuleManager.h"
#include "Primitives/AEC/InworldAECTypes.h"

#include <Interfaces/IPluginManager.h>

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

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

	// Create AEC configuration
	inworld::LocalAECFilterConfig local_aec_config;
	inworld::Convert(FInworldAECCreationConfig{}, local_aec_config);

	// Create AEC instance
	if (const inworld::StatusOr_AECFilterInterface status_or_aec{inworld::AECFilterFactory().CreateAECFilter(local_aec_config)};
		status_or_aec.ok())
	{
		RETURN_SUCCESS(MakeShared<inworld::AECFilterInterface>(status_or_aec.value()))
	}
	else
	{
		RETURN_ERROR(
			FString::Format(TEXT("{0}: {1}"), {*FString(__func__), UTF8_TO_TCHAR(status_or_aec.status().ToString().c_str())}))
	}
#undef RETURN_ERROR
#undef RETURN_SUCCESS
}

namespace inworld
{
namespace Tasks
{
namespace AEC
{

UE::Tasks::TTask<TVariant<TSharedPtr<inworld::AECFilterInterface>, FString>> CreateAECTask(
	const FInworldAECCreationConfig& AECConfig)
{
	return UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[]()
		{
			return CreateAEC();
		});
}

UE::Tasks::TTask<TVariant<FInworldData_Audio, FString>> CreateFilterAudioTask(TSharedPtr<inworld::AECFilterInterface> AECInterface,
	const FInworldAECConfig& Config, const FInworldData_Audio& MicrophoneAudio, const FInworldData_Audio& ReferenceAudio)
{
#define RETURN_ERROR(Error) return TVariant<FInworldData_Audio, FString>{TInPlaceType<FString>(), Error};
#define RETURN_SUCCESS(Object) return TVariant<FInworldData_Audio, FString>{TInPlaceType<FInworldData_Audio>(), Object};

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

			// Create microphone audio chunk
			inworld::AudioChunk mic_audio_chunk;
			mic_audio_chunk.sample_rate(MicrophoneAudio.SampleRate);
			inworld::vector_float mic_vector_float = inworld::memcpy_to_vector_float(
				(float*) MicrophoneAudio.Waveform->GetData(), MicrophoneAudio.Waveform->Num() / sizeof(float));
			mic_audio_chunk.data(&mic_vector_float);

			// Create reference audio chunk (speaker audio)
			inworld::AudioChunk ref_audio_chunk;
			ref_audio_chunk.sample_rate(ReferenceAudio.SampleRate);
			inworld::vector_float ref_vector_float = inworld::memcpy_to_vector_float(
				(float*) ReferenceAudio.Waveform->GetData(), ReferenceAudio.Waveform->Num() / sizeof(float));
			ref_audio_chunk.data(&ref_vector_float);

			// Call AEC filtering
			if (const inworld::StatusOr_AudioChunk status_or_filtered_audio{
					AECInterface->FilterAudio(mic_audio_chunk, ref_audio_chunk)};
				status_or_filtered_audio.ok())
			{
				// Get filtered audio
				const inworld::AudioChunk filtered_audio_chunk = status_or_filtered_audio.value();

				// Convert back to UE audio format
				FInworldData_Audio filtered_audio;
				filtered_audio.SampleRate = filtered_audio_chunk.sample_rate();

				// Get audio data
				inworld::vector_float* filtered_vector = filtered_audio_chunk.data();

				int32 sample_count = filtered_vector->size();

				// Convert to UE audio data format
				filtered_audio.Waveform = MakeShared<TArray<uint8>>();
				filtered_audio.Waveform->SetNumUninitialized(sample_count * sizeof(float));

				// Copy audio data
				for (int32 i = 0; i < sample_count; ++i)
				{
					float sample = filtered_vector->get(i);
					FMemory::Memcpy(filtered_audio.Waveform->GetData() + (i * sizeof(float)), &sample, sizeof(float));
				}

				RETURN_SUCCESS(filtered_audio)
			}
			else
			{
				RETURN_ERROR(FString::Format(
					TEXT("{0}: {1}"), {*FString(__func__), UTF8_TO_TCHAR(status_or_filtered_audio.status().ToString().c_str())}))
			}
		});

#undef RETURN_ERROR
#undef RETURN_SUCCESS
}

}	 // namespace AEC
}	 // namespace Tasks
}	 // namespace inworld
