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

#include "Audio/InworldVoiceAudioComponent.h"

#include "Async/Async.h"
#include "Async/TaskGraphInterfaces.h"
#include "AudioDevice.h"
#include "Primitives/AEC/InworldAECTypes.h"
#include "Sound/SoundWave.h"
#include "Sound/SoundWaveProcedural.h"
#include "TimerManager.h"

#include <Engine/World.h>
#include <GameFramework/Actor.h>

void UInworldVoiceAudioComponent::BeginPlay()
{
	Super::BeginPlay();

	SoundStreaming = NewObject<USoundWaveProcedural>();

	SoundStreaming->OnSoundWaveProceduralUnderflow =
		FOnSoundWaveProceduralUnderflow::CreateUObject(this, &UInworldVoiceAudioComponent::GenerateData);

	SetSound(SoundStreaming);
}

void UInworldVoiceAudioComponent::EndPlay(EEndPlayReason::Type Reason)
{
	Super::EndPlay(Reason);

	SoundStreaming->OnSoundWaveProceduralUnderflow = {};
	SoundStreaming = nullptr;

	SetSound(nullptr);
}

void UInworldVoiceAudioComponent::QueueVoice(const FInworldData_TTSOutput& Voice)
{
	FScopeLock ScopeLock(&QueueLock);

	if (!Voice.Audio.Waveform.IsValid() || Voice.Audio.Waveform->IsEmpty())
	{
		return;
	}

	if (!IsValid(SoundStreaming))
	{
		return;
	}

	if (Voices.Num() == 0)
	{
		SourceSampleRate = Voice.Audio.SampleRate;
		SoundStreaming->SetSampleRate(Voice.Audio.SampleRate);
		SoundStreaming->NumChannels = 1;
	}

	Voices.Add(Voice);

	if (bIsPaused)
	{
		SetPaused(false);
	}

	if (!IsPlaying())
	{
		Play();
	}
}

void UInworldVoiceAudioComponent::Interrupt()
{
	FScopeLock ScopeLock(&QueueLock);
	Stop();
	if (Voices.Num() > 0)
	{
		TArray<FInworldData_TTSOutput> InterruptedVoices{Voices};

		TWeakObjectPtr<UInworldVoiceAudioComponent> WeakSelf = this;
		AsyncTask(ENamedThreads::GameThread,
			[WeakSelf, InterruptedVoices]()
			{
				if (!WeakSelf.IsValid())
				{
					return;
				}
				if (!WeakSelf->GetOwner()->IsPendingKillPending())
				{
					for (const FInworldData_TTSOutput& Voice : InterruptedVoices)
					{
						WeakSelf->OnVoiceAudioInterrupt.Broadcast(WeakSelf.Get(), Voice);
						WeakSelf->OnVoiceAudioInterruptNative.Broadcast(WeakSelf.Get(), Voice);
					}
				}
			});
	}

	Voices = {};
	NumSoundDataBytesPlayed = 0;
}

void UInworldVoiceAudioComponent::GenerateData(USoundWaveProcedural* InProceduralWave, int32 SamplesRequired)
{
	FScopeLock ScopeLock(&QueueLock);
	if (Voices.Num() == 0)
	{
		return;
	}
	if (NumSoundDataBytesPlayed == Voices[0].Audio.Waveform->Num())
	{
		const FInworldData_TTSOutput CompletedVoice{Voices[0]};
		Voices.RemoveAt(0);
		NumSoundDataBytesPlayed = 0;

		TWeakObjectPtr<UInworldVoiceAudioComponent> WeakSelf = this;
		AsyncTask(ENamedThreads::GameThread,
			[WeakSelf, CompletedVoice]()
			{
				if (!WeakSelf.IsValid())
				{
					return;
				}
				if (!WeakSelf->GetOwner()->IsPendingKillPending())
				{
					FScopeLock ScopeLock(&WeakSelf->QueueLock);
					WeakSelf->OnVoiceAudioComplete.Broadcast(WeakSelf.Get(), CompletedVoice);
					WeakSelf->OnVoiceAudioCompleteNative.Broadcast(WeakSelf.Get(), CompletedVoice);
					if (WeakSelf->Voices.Num() == 0)
					{
						WeakSelf->SoundStreaming->ResetAudio();
						WeakSelf->Stop();
					}
					else
					{
						WeakSelf->Play();
					}
				}
			});
	}
	else
	{
		const FInworldData_TTSOutput& StartedVoice{Voices[0]};
		if (NumSoundDataBytesPlayed == 0)
		{
			SoundStreaming->SetSampleRate(StartedVoice.Audio.SampleRate);
			SoundStreaming->NumChannels = 1;

			TWeakObjectPtr<UInworldVoiceAudioComponent> WeakSelf = this;
			AsyncTask(ENamedThreads::GameThread,
				[WeakSelf, StartedVoice]()
				{
					if (!WeakSelf.IsValid())
					{
						return;
					}
					if (!WeakSelf->GetOwner()->IsPendingKillPending())
					{
						// Start voice
						FScopeLock ScopeLock(&WeakSelf->QueueLock);
						WeakSelf->OnVoiceAudioStart.Broadcast(WeakSelf.Get(), StartedVoice);
						WeakSelf->OnVoiceAudioStartNative.Broadcast(WeakSelf.Get(), StartedVoice);
					}
				});
		}

		const int NextSampleCount = FMath::Min(SamplesRequired, (Voices[0].Audio.Waveform->Num() - NumSoundDataBytesPlayed) / 2);

		TArray<float> SpeechDataFloat;
		SpeechDataFloat.SetNumUninitialized(NextSampleCount / 2);
		FMemory::Memcpy(
			SpeechDataFloat.GetData(), Voices[0].Audio.Waveform->GetData() + NumSoundDataBytesPlayed, NextSampleCount * 2);
		TArray<int16> SpeechData16;
		for (const float& SpeechData : SpeechDataFloat)
		{
			SpeechData16.Add(SpeechData * 32767);
		}
		TArray<uint8> AudioData{(const uint8*) SpeechData16.GetData(), SpeechData16.Num() * 2};
		InProceduralWave->QueueAudio(AudioData.GetData(), NextSampleCount);
		NumSoundDataBytesPlayed += NextSampleCount * 2;

		// for lipsync viseme generation. We pass the SpeechData16 so it
		OnSignedAudioDataGenerated.ExecuteIfBound(MoveTemp(SpeechData16), SourceSampleRate);

		// On playback
		{
			float BytesPerSecond = 1.f * (32.f /*bps*/ / 8.f) * (float) StartedVoice.Audio.SampleRate;
			float PlaybackDuration = (float) StartedVoice.Audio.Waveform->Num() / BytesPerSecond;
			float PlaybackTimeSeconds = NumSoundDataBytesPlayed / BytesPerSecond;

			TWeakObjectPtr<UInworldVoiceAudioComponent> WeakSelf = this;
			AsyncTask(ENamedThreads::GameThread,
				[WeakSelf, PlaybackDuration, PlaybackTimeSeconds]()
				{
					if (!WeakSelf.IsValid())
					{
						return;
					}
					if (!WeakSelf->GetOwner()->IsPendingKillPending())
					{
						FScopeLock ScopeLock(&WeakSelf->QueueLock);
						WeakSelf->OnVoiceAudioPlayback.Broadcast(WeakSelf.Get(), PlaybackDuration, PlaybackTimeSeconds);
						WeakSelf->OnVoiceAudioPlaybackNative.Broadcast(WeakSelf.Get(), PlaybackDuration, PlaybackTimeSeconds);
					}
				});
		}
	}
}
