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

#include "Components/InworldBasePlayerComponent.h"

#include "Audio/InworldAudioDumpUtils.h"
#include "Conversation/InworldConversationTarget.h"
#include "Engine/World.h"
#include "Graph/InworldGraph.h"
#include "Graph/Nodes/InworldNode_STT.h"
#include "HAL/FileManager.h"
#include "InworldCharacter.h"
#include "TimerManager.h"

void UInworldBasePlayerComponent::SetConversationTarget(TScriptInterface<IInworldConversationTarget> NewConversationTarget)
{
	if (ConversationTarget == NewConversationTarget)
	{
		return;
	}
	ConversationTarget = NewConversationTarget;

	EvaluateTalking();

	OnConversationTargetUpdate.Broadcast(this, ConversationTarget);
	OnConversationTargetUpdateNative.Broadcast(this, ConversationTarget);
}

void UInworldBasePlayerComponent::RemoveConversationTarget()
{
	SetConversationTarget(nullptr);
}

TScriptInterface<IInworldConversationTarget> UInworldBasePlayerComponent::GetConversationTarget()
{
	return ConversationTarget;
}

void UInworldBasePlayerComponent::SendMessage(const FInworldDataHandle& Message)
{
	if (IsValid(ConversationTarget.GetObject()))
	{
		IInworldConversationTarget::Execute_SendMessage(ConversationTarget.GetObject(), Message);
		OnPlayerMessage.Broadcast(this, ConversationTarget, Message);
		OnPlayerMessageNative.Broadcast(this, ConversationTarget, Message);
	}
	else
	{
		UE_LOG(LogInworldCharacter, Warning,
			TEXT("UInworldBasePlayerComponent::SendMessage - failed to send message."
				 " Invalid conversation target."))
	}
}

void UInworldBasePlayerComponent::SetMicrophone(bool Enabled)
{
	bMicEnabled = Enabled;
	EvaluateTalking();
}

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

	UInworldAudioCapture::CreateInworldAudioCapture(this, AudioCaptureConfig,
		FOnInworldAudioCaptureCreatedNative::CreateUObject(this, &UInworldBasePlayerComponent::OnAudioCaptureCreated));
#if WITH_EDITOR && !UE_BUILD_SHIPPING
	bDebugDumpAudio = inworld::audio::debug::GetCVarDebugDump().GetValueOnGameThread();
	if (bDebugDumpAudio)
	{
		DebugDumpPlayerUtteranceAudioFile = FPaths::Combine(inworld::audio::debug::GetCVarDebugDumpPath().GetValueOnGameThread(),
			FString::Format(TEXT("PlayerUtteranceAudio_{0}.wav"), {GetName()}));
	}
#endif

	STTGraph = NewObject<UInworldGraph>(this);
	STTGraph->SetGraphId(TEXT("STT"));
	STTGraph->Nodes.Add(UInworldNode_STT::CreateNative(this, TEXT("STT"), {true, "", StreamingSTTComponentId}));
	STTGraph->StartNodes.Add(TEXT("STT"));
	STTGraph->EndNodes.Add(TEXT("STT"));

	STTGraph->GetGraphInstance(FOnGraphCompiledNative::CreateWeakLambda(this,
		[this](UInworldGraph* CompiledGraph, bool bSuccess)
		{
			if (bSuccess)
			{
				TrySignalReady();
			}
			else
			{
				UE_LOG(LogInworldCharacter, Error, TEXT("UInworldBasePlayerComponent::BeginPlay - Failed to compile STT graph."));
			}
		}));
}

void UInworldBasePlayerComponent::OnAudioCaptureCreated(UInworldAudioCapture* AudioCapture)
{
	InworldAudioCapture = AudioCapture;
	NumVADForgivenessFrames = VADBufferTime * AudioCapture->CaptureFrequency;
	InworldAudioCapture->OnAudioCaptureNative.AddUObject(this, &UInworldBasePlayerComponent::OnAudioCapture);
	if (bStartMicEnabled)
	{
		SetMicrophone(true);
	}
	TrySignalReady();
}

void UInworldBasePlayerComponent::OnAudioCapture(
	const UInworldAudioCapture* AudioCapture, const FInworldData_Audio& AudioData, bool VAD)
{
	OnPlayerAudioCapture.Broadcast(this, AudioData, VAD);
	OnPlayerAudioCaptureNative.Broadcast(this, AudioData, VAD);
	ProcessAudioData(AudioData, VAD);
}

void UInworldBasePlayerComponent::EvaluateTalking()
{
	if (!IsValid(InworldAudioCapture))
	{
		bTalking = false;
		OnFinishedSpeaking();
		return;
	}

	const bool bHasConversationTarget = IsValid(ConversationTarget.GetObject());
	const bool bCanTalk = bMicEnabled && bHasConversationTarget;

	if (!bTalking && bCanTalk)
	{
		bTalking = true;
		InworldAudioCapture->StartCapture();
	}
	else if (bTalking && !bCanTalk)
	{
		bTalking = false;
		InworldAudioCapture->StopCapture();
		PrefixAudioDataFrames.Empty();
		OnFinishedSpeaking();
	}
}

void UInworldBasePlayerComponent::ProcessAudioData(const FInworldData_Audio& AudioData, bool VAD)
{
	if (VAD)
	{
		CurrentVADForgivenessFrame = 0;
		if (!InworldAudioStream.IsValid() && IsValid(ConversationTarget.GetObject()))
		{
			InworldAudioStream = MakeShared<FInworldData_DataStream_AudioChunk>();
#if WITH_EDITOR && !UE_BUILD_SHIPPING
			if (bDebugDumpAudio)
			{
				InworldAudioStream->SetDebugDumpFileName(DebugDumpPlayerUtteranceAudioFile);
			}
#endif
			for (const auto& AudioDataFrame : PrefixAudioDataFrames)
			{
				InworldAudioStream->Write(AudioDataFrame);
			}

#if WITH_EDITOR && !UE_BUILD_SHIPPING
			if (bDebugDumpAudio)
			{
				IFileManager::Get().Delete(*DebugDumpPlayerUtteranceAudioFile);
			}
#endif

			STTGraph->Execute(this, *InworldAudioStream.Get(),
				FOnGraphResultNative::CreateLambda(
					[WeakThis = TWeakObjectPtr<UInworldBasePlayerComponent>(this)](
						const FString& ExecutionId, const FString& NodeId, const FInworldDataHandle& Data)
					{
						if (!Data.IsValid() || !Data.IsA<FInworldData_Text>())
						{
							UE_LOG(LogInworldCharacter, Error,
								TEXT("UInworldBasePlayerComponent::ProcessAudioData - "
									 "Failed to convert speech to text."));
							return;
						}

						if (FInworldData_Text* TextData = Data.Unwrap<FInworldData_Text>(); TextData->Text.IsEmpty())
						{
							// Do not send empty text messages
							return;
						}

						if (WeakThis.IsValid())
						{
							WeakThis->SendMessage(Data);
						}
						else
						{
							UE_LOG(LogInworldCharacter, Error,
								TEXT("UInworldBasePlayerComponent::ProcessAudioData - "
									 "Failed to send STT result. Invalid target."));
						}
					}));
		}
		InworldAudioStream->Write(AudioData);
	}
	else if (InworldAudioStream.IsValid())
	{
		if (CurrentVADForgivenessFrame < NumVADForgivenessFrames)
		{
			InworldAudioStream->Write(AudioData);
			CurrentVADForgivenessFrame++;
		}
		if (CurrentVADForgivenessFrame == NumVADForgivenessFrames)
		{
			OnFinishedSpeaking();
		}
	}
	else
	{
		PrefixAudioDataFrames.AddTail(AudioData);
		while (PrefixAudioDataFrames.Num() > NumVADForgivenessFrames)
		{
			PrefixAudioDataFrames.RemoveNode(PrefixAudioDataFrames.GetHead());
		}
	}
}

void UInworldBasePlayerComponent::OnFinishedSpeaking()
{
	if (InworldAudioStream.IsValid())
	{
		InworldAudioStream->Close();
		InworldAudioStream.Reset();
		CurrentVADForgivenessFrame = 0;
	}
}

void UInworldBasePlayerComponent::TrySignalReady()
{
	if (bIsReady)
	{
		return;
	}

	if (InworldAudioCapture != nullptr && STTGraph != nullptr && STTGraph->IsCompiled())
	{
		bIsReady = true;
		OnPlayerReady.Broadcast(this);
		OnPlayerReadyNative.Broadcast(this);
	}
}
