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

#include "Audio/InworldAudioCapture.h"

#include "Algo/Accumulate.h"
#include "Async/Async.h"
#include "Audio/InworldAudioCapture_MasterSubmix.h"
#include "Audio/InworldAudioCapture_Microphone.h"
#include "Audio/InworldAudioDumpUtils.h"
#include "AudioResampler.h"
#include "HAL/FileManager.h"
#include "HAL/PlatformFileManager.h"
#include "InworldRuntime.h"
#include "InworldRuntimeSubsystem.h"
#include "Misc/ScopeTryLock.h"
#include "Primitives/AEC/InworldAECTypes.h"
#include "Primitives/VAD/InworldVADTypes.h"
#include "Runtime/Launch/Resources/Version.h"
#if WITH_EDITOR
#include "Editor.h"
#else
#include "Engine/Engine.h"
#endif	  // WITH_EDITOR

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

namespace inworld
{
namespace audio
{
void ConvertAudioCaptureDataToMono(FInworldAudioCaptureData& InAudioCaptureData)
{
	if (InAudioCaptureData.NumChannels != 1)
	{
		int32 DataOffset{0};
		const int32 NumFrames{InAudioCaptureData.AudioData.Num() / InAudioCaptureData.NumChannels};
		for (int32 CurrentFrame = 0; CurrentFrame < NumFrames; CurrentFrame++)
		{
			InAudioCaptureData.AudioData[CurrentFrame] =
				Algo::Accumulate(
					TArray<float>{InAudioCaptureData.AudioData.GetData() + DataOffset, InAudioCaptureData.NumChannels}, 0.f) /
				InAudioCaptureData.NumChannels;
			DataOffset += InAudioCaptureData.NumChannels;
		}
		// Resize array to match mono output
		InAudioCaptureData.AudioData.SetNum(NumFrames);
		InAudioCaptureData.NumChannels = 1;
	}
}
}	 // namespace audio
}	 // namespace inworld

void UInworldAudioCapture::CreateInworldAudioCapture(
	UObject* WorldContextObject, const FInworldAudioCaptureConfig& Config, FOnInworldAudioCaptureCreated Callback)
{
	FOnInworldAudioCaptureCreatedNative NativeCallback;
	if (Callback.IsBound())
	{
		NativeCallback.BindUFunction(Callback.GetUObject(), Callback.GetFunctionName());
	}
	CreateInworldAudioCapture(WorldContextObject, Config, NativeCallback);
}

void UInworldAudioCapture::CreateInworldAudioCapture(
	UObject* WorldContextObject, const FInworldAudioCaptureConfig& Config, FOnInworldAudioCaptureCreatedNative Callback)
{
	AsyncTask(ENamedThreads::AnyThread,
		[WorldContext = TWeakObjectPtr<UObject>(WorldContextObject), Config, Callback]()
		{
			const bool bAECValid = Config.bEnableAEC && UInworldRuntimeSubsystem::Get()->GetAEC().is_valid();
			const bool bVADValid = Config.bEnableVAD && UInworldRuntimeSubsystem::Get()->GetVAD().is_valid();
			AsyncTask(ENamedThreads::GameThread,
				[WorldContext, PerSecondSendRate = Config.CaptureFrequency, bAECValid, AECConfig = Config.AECConfig, bVADValid,
					VADConfig = Config.VADConfig, Callback]()
				{
					UInworldAudioCapture* const AudioCapture =
						NewObject<UInworldAudioCapture>(WorldContext.IsValid() ? WorldContext.Get() : nullptr);
					AudioCapture->CaptureFrequency = PerSecondSendRate;
					AudioCapture->bUseAEC = bAECValid;
					AudioCapture->AECConfig = AECConfig;
					AudioCapture->bUseVAD = bVADValid;
					AudioCapture->VADConfig = VADConfig;
					Callback.ExecuteIfBound(AudioCapture);
				});
		});
}

void UInworldAudioCapture::PostInitProperties()
{
	Super::PostInitProperties();

	if (HasAnyFlags(RF_ClassDefaultObject))
	{
		return;
	}

#if WITH_EDITOR && !UE_BUILD_SHIPPING
	bDebugDumpAudio = inworld::audio::debug::GetCVarDebugDump().GetValueOnGameThread();
	if (bDebugDumpAudio)
	{
		const FString DebugDumpAudioPath{inworld::audio::debug::GetCVarDebugDumpPath().GetValueOnGameThread()};
		DebugDumpInputAudioFile = FPaths::Combine(DebugDumpAudioPath, FString::Format(TEXT("InputAudio_{0}.wav"), {GetName()}));
		IFileManager::Get().Delete(*DebugDumpInputAudioFile);
		DebugDumpOutputAudioFile = FPaths::Combine(DebugDumpAudioPath, FString::Format(TEXT("OutputAudio_{0}.wav"), {GetName()}));
		IFileManager::Get().Delete(*DebugDumpOutputAudioFile);
		DebugDumpResultAudioFile = FPaths::Combine(DebugDumpAudioPath, FString::Format(TEXT("ResultAudio_{0}.wav"), {GetName()}));
		IFileManager::Get().Delete(*DebugDumpResultAudioFile);
	}
#endif

	InputAudioCapture = MakeShared<FInworldAudioCapture_Microphone>(this,
		[WeakThis = TWeakObjectPtr<UInworldAudioCapture>(this)](FInworldAudioCaptureData AudioCaptureData)
		{
			if (!WeakThis.IsValid() || !WeakThis->bCapturingAudio)
			{
				return;
			}
			inworld::audio::ConvertAudioCaptureDataToMono(AudioCaptureData);
			{
				WeakThis->InputBuffer.PCQueue.Enqueue(AudioCaptureData.AudioData);
				WeakThis->InputBuffer.SampleRate = AudioCaptureData.SampleRate;
			}
			WeakThis->TryProcessAudio();
		});

	OutputAudioCapture = MakeShared<FInworldAudioCapture_MasterSubmix>(this,
		[WeakThis = TWeakObjectPtr<UInworldAudioCapture>(this)](FInworldAudioCaptureData AudioCaptureData)
		{
			if (!WeakThis.IsValid() || !WeakThis->bCapturingAudio)
			{
				return;
			}
			inworld::audio::ConvertAudioCaptureDataToMono(AudioCaptureData);
			{
				WeakThis->OutputBuffer.PCQueue.Enqueue(AudioCaptureData.AudioData);
				WeakThis->OutputBuffer.SampleRate = AudioCaptureData.SampleRate;
			}
			WeakThis->TryProcessAudio();
		});
}

void UInworldAudioCapture::TryProcessAudio()
{
	FScopeTryLock ProcessAudoScopedLock(&ProcessAudioCriticalSection);
	if (!ProcessAudoScopedLock.IsLocked())
	{
		return;
	}

	TArray<float> InputChunk;
	while (InputBuffer.PCQueue.Dequeue(InputChunk))
	{
		InputBuffer.Data.Append(InputChunk);
	}
	// mult by 2 from Buffer (uint8) to PCM (uint16)
	const int32 InputSampleSendSize = (InputBuffer.SampleRate / CaptureFrequency) * 2;
	if (InputBuffer.Data.Num() <= InputSampleSendSize)
	{
		return;
	}

	TArray<float> OutputChunk;
	while (OutputBuffer.PCQueue.Dequeue(OutputChunk))
	{
		OutputBuffer.Data.Append(OutputChunk);
	}
	// mult by 2 from Buffer (uint8) to PCM (uint16)
	const int32 OutputSampleSendSize = (OutputBuffer.SampleRate / CaptureFrequency) * 2;
	if (OutputBuffer.Data.Num() <= OutputSampleSendSize)
	{
		return;
	}

	FInworldData_Audio InputAudio;
	InputAudio.Waveform = MakeShared<TArray<uint8>>();
	InputAudio.Waveform->SetNumUninitialized(InputSampleSendSize * sizeof(float));
	InputAudio.SampleRate = InputBuffer.SampleRate;
	FMemory::Memcpy(InputAudio.Waveform->GetData(), InputBuffer.Data.GetData(), InputSampleSendSize * sizeof(float));
	FMemory::Memcpy(InputBuffer.Data.GetData(), InputBuffer.Data.GetData() + InputSampleSendSize,
		(InputBuffer.Data.Num() - InputSampleSendSize) * sizeof(float));
	InputBuffer.Data.SetNum(InputBuffer.Data.Num() - InputSampleSendSize);

	FInworldData_Audio OutputAudio;
	OutputAudio.Waveform = MakeShared<TArray<uint8>>();
	OutputAudio.Waveform->SetNumUninitialized(OutputSampleSendSize * sizeof(float));
	OutputAudio.SampleRate = OutputBuffer.SampleRate;
	FMemory::Memcpy(OutputAudio.Waveform->GetData(), OutputBuffer.Data.GetData(), OutputSampleSendSize * sizeof(float));
	FMemory::Memcpy(OutputBuffer.Data.GetData(), OutputBuffer.Data.GetData() + OutputSampleSendSize,
		(OutputBuffer.Data.Num() - OutputSampleSendSize) * sizeof(float));
	OutputBuffer.Data.SetNum(OutputBuffer.Data.Num() - OutputSampleSendSize);

#if WITH_EDITOR && !UE_BUILD_SHIPPING
	if (bDebugDumpAudio)
	{
		inworld::audio::debug::DumpAudio(DebugDumpInputAudioFile, InputAudio);
		inworld::audio::debug::DumpAudio(DebugDumpOutputAudioFile, OutputAudio);
	}
#endif

	auto AECTask = UE::Tasks::Launch(
		UE_SOURCE_LOCATION,
		[bUseAEC = bUseAEC, AECConfig = AECConfig, InputAudio = MoveTemp(InputAudio), OutputAudio = MoveTemp(OutputAudio)]()
		{
			static inworld::AECFilterInterface null_eac{inworld::AECFilterInterface::null()};
			inworld::AECFilterInterface& aec = bUseAEC ? UInworldRuntimeSubsystem::Get()->GetAEC() : null_eac;
			if (aec.is_valid())
			{
				// Create microphone audio chunk
				inworld::AudioChunk mic_audio_chunk;
				mic_audio_chunk.sample_rate(InputAudio.SampleRate);
				inworld::vector_float mic_vector_float = inworld::memcpy_to_vector_float(
					(float*) InputAudio.Waveform->GetData(), InputAudio.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(OutputAudio.SampleRate);
				inworld::vector_float ref_vector_float = inworld::memcpy_to_vector_float(
					(float*) OutputAudio.Waveform->GetData(), OutputAudio.Waveform->Num() / sizeof(float));
				ref_audio_chunk.data(&ref_vector_float);

				// Call AEC filtering
				if (const inworld::StatusOr_AudioChunk status_or_filtered_audio{aec.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 FilteredAudio;
					FilteredAudio.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
					FilteredAudio.Waveform = MakeShared<TArray<uint8>>();
					FilteredAudio.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(FilteredAudio.Waveform->GetData() + (i * sizeof(float)), &sample, sizeof(float));
					}

					return FilteredAudio;
				}
				else
				{
					UE_LOG(LogInworldRuntime, Error, TEXT("Failed to filter audio: %s"),
						UTF8_TO_TCHAR(status_or_filtered_audio.status().ToString().c_str()));
				}
			}

			return InputAudio;
		},
		ProcessAudioTask);

	auto VADTask = UE::Tasks::Launch(
		UE_SOURCE_LOCATION,
		[AECTask, bUseVAD = bUseVAD, VADConfig = VADConfig]() mutable
		{
			const FInworldData_Audio& AudioData{AECTask.GetResult()};
			bool bVAD = false;

			static inworld::VADInterface null_vad{inworld::VADInterface::null()};
			inworld::VADInterface& vad = bUseVAD ? UInworldRuntimeSubsystem::Get()->GetVAD() : null_vad;
			if (vad.is_valid())
			{
				inworld::VoiceActivityDetectionConfig voice_activity_detection_config;
				voice_activity_detection_config.speech_threshold(VADConfig.SpeechThreshold);

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

				if (const inworld::StatusOr_int status_or_vad_result{
						vad.DetectVoiceActivity(audio_chunk, voice_activity_detection_config)};
					status_or_vad_result.ok())
				{
					bVAD = {status_or_vad_result.value() != INDEX_NONE};
				}
				else
				{
					UE_LOG(LogInworldRuntime, Error, TEXT("Failed to vad audio: %s"),
						UTF8_TO_TCHAR(status_or_vad_result.status().ToString().c_str()));
				}
			}
			return TTuple<FInworldData_Audio, bool>(AudioData, bVAD);
		},
		AECTask);

	ProcessAudioTask = UE::Tasks::Launch(
		UE_SOURCE_LOCATION,
		[VADTask, WeakThis = TWeakObjectPtr<UInworldAudioCapture>(this)]() mutable
		{
			if (!WeakThis.IsValid())
			{
				return;
			}

			const auto& Result{VADTask.GetResult()};
			const FInworldData_Audio ResultAudio = Result.Key;
			const bool bVAD = Result.Value;

#if WITH_EDITOR && !UE_BUILD_SHIPPING
			if (WeakThis->bDebugDumpAudio)
			{
				inworld::audio::debug::DumpAudio(WeakThis->DebugDumpResultAudioFile, ResultAudio);
			}
#endif
			AsyncTask(ENamedThreads::GameThread,
				[WeakThis, ResultAudio, bVAD]()
				{
					if (!WeakThis.IsValid())
					{
						return;
					}

					WeakThis->OnAudioCapture.Broadcast(WeakThis.Get(), ResultAudio, bVAD);
					WeakThis->OnAudioCaptureNative.Broadcast(WeakThis.Get(), ResultAudio, bVAD);
				});
		},
		VADTask);
}

void UInworldAudioCapture::BeginDestroy()
{
	StopCapture();

	Super::BeginDestroy();
}

void UInworldAudioCapture::StartCapture()
{
	if (!IsInAudioThread())
	{
		FAudioThread::RunCommandOnAudioThread(
			[WeakThis = TWeakObjectPtr<UInworldAudioCapture>(this)]()
			{
				if (WeakThis.IsValid())
				{
					WeakThis->StartCapture();
				}
			});
		return;
	}

	if (bCapturingAudio)
	{
		return;
	}
	if (!InputAudioCapture.IsValid())
	{
		return;
	}
	InputAudioCapture->StartCapture();

	if (OutputAudioCapture.IsValid())
	{
		OutputAudioCapture->StartCapture();
	}
	bCapturingAudio = true;
}

void UInworldAudioCapture::StopCapture()
{
	if (!IsInAudioThread())
	{
		FAudioThread::RunCommandOnAudioThread(
			[WeakThis = TWeakObjectPtr<UInworldAudioCapture>(this)]()
			{
				if (WeakThis.IsValid())
				{
					WeakThis->StopCapture();
				}
			});
		return;
	}

	if (!bCapturingAudio)
	{
		return;
	}
	if (InputAudioCapture.IsValid())
	{
		InputAudioCapture->StopCapture();
	}
	if (OutputAudioCapture.IsValid())
	{
		OutputAudioCapture->StopCapture();
	}

	bCapturingAudio = false;

	InputBuffer.Data.Empty();
	OutputBuffer.Data.Empty();
	FScopeLock ProcessAudoScopedLock(&ProcessAudioCriticalSection);
	ProcessAudioTask = {};
}
