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

#include "InworldLipsyncComponent.h"
#include "InworldLipsync.h"
#include "GameFramework/Actor.h"
#include "Interfaces/IPluginManager.h"
#include "Misc/Paths.h"
#include "Async/Async.h"
#include "Async/TaskGraphInterfaces.h"
#include "Graph/InworldData.h"
#include "Audio/InworldVoiceAudioComponent.h"

DECLARE_STATS_GROUP(TEXT("InworldLipsync"), STATGROUP_InworldLipsync, STATCAT_Advanced);
DECLARE_CYCLE_STAT(TEXT("Lipsync QueueAudio"), STAT_InworldLipsync_QueueAudio, STATGROUP_InworldLipsync);
DECLARE_CYCLE_STAT(TEXT("Lipsync Resample Waveform"), STAT_InworldLipsync_DownsampleWaveform, STATGROUP_InworldLipsync);
DECLARE_CYCLE_STAT(TEXT("Lipsync Sphinx Process"), STAT_InworldLipsync_SphinxProcess, STATGROUP_InworldLipsync);
DECLARE_CYCLE_STAT(TEXT("Lipsync Sphinx Hyp"), STAT_InworldLipsync_SphinxHyp, STATGROUP_InworldLipsync);

FAudioProcessingRunnable::FAudioProcessingRunnable(UInworldLipsyncComponent* InOwnerComponent)
	: OwnerComponent(InOwnerComponent), bStopping(false), bVoiceComplete(false), Thread(nullptr)

{
	// Start the thread
	Thread = FRunnableThread::Create(this, TEXT("InworldLipsyncProcessingThread"), 0, TPri_Normal);
	UE_LOG(LogInworldLipsync, Log, TEXT("Lipsync Audio Processing Thread created."));
}

FAudioProcessingRunnable::~FAudioProcessingRunnable()
{
	EnsureCompletion();
	if (Thread)
	{
		delete Thread;
		Thread = nullptr;
	}
}

void FAudioProcessingRunnable::InitializePocketSphinx()
{
	// --- Get paths to our model files in the Content folder ---
	FString BaseDir = IPluginManager::Get().FindPlugin("InworldLipsync")->GetBaseDir();
	FString ModelPath = FPaths::Combine(*BaseDir, TEXT("Model"));

	// Define paths using the robust base path
	FString HmmPath = FPaths::Combine(*ModelPath, TEXT("en-us"));
	FString PhoneLmPath = FPaths::Combine(*ModelPath, TEXT("en-us-phone.lm.bin"));
	FString DictPath = FPaths::Combine(*ModelPath, TEXT("cmudict-en-us.dict"));
	FString SphinxLogPath = FPaths::ProjectLogDir() + TEXT("pocketsphinx.log");

	// --- Configure PocketSphinx ---
	// We use -allphone to get phonemes instead of words.
	// We use -logfn to redirect log spam away from the Unreal console.
	SphinxConfig = ps_config_init(NULL);
	ps_config_set_int(SphinxConfig, "samprate", PocketSphinxSampleRate);
	ps_config_set_int(SphinxConfig, "maxwpf", 40);
	ps_config_set_str(SphinxConfig, "hmm", TCHAR_TO_ANSI(*HmmPath));
	ps_config_set_str(SphinxConfig, "dict", TCHAR_TO_ANSI(*DictPath));
	ps_config_set_str(SphinxConfig, "logfn", TCHAR_TO_ANSI(*SphinxLogPath));	// Use "/dev/null" on Linux/Mac
	ps_config_set_bool(SphinxConfig, "bestpath", true);

	ps_config_set_float(SphinxConfig, "lw", 2.0);	 // Language weight, helps with phoneme separation
	ps_config_set_str(SphinxConfig, "allphone", TCHAR_TO_ANSI(*PhoneLmPath));

	// --- Initialize the decoder ---
	SphinxDecoder = ps_init(SphinxConfig);
	if (!SphinxDecoder)
	{
		UE_LOG(LogInworldLipsync, Error, TEXT("PocketSphinx failed to initialize."));
	}
}

void FAudioProcessingRunnable::ShutdownPocketSphinx()
{
	if (SphinxDecoder)
	{
		ps_free(SphinxDecoder);
		SphinxDecoder = nullptr;
	}
	if (SphinxConfig)
	{
		ps_config_free(SphinxConfig);
		SphinxConfig = nullptr;
	}
}

uint32 FAudioProcessingRunnable::Run()
{
	// Main processing loop. This runs until Stop() is called.
	if (!bIsPocketSphinxCreated)
	{
		bIsPocketSphinxCreated = true;
		UE_LOG(LogInworldLipsync, Log, TEXT("Creating PocketSphinxDecoder"));
		InitializePocketSphinx();
		if (SphinxDecoder == nullptr)
		{
			bStopping = true;
		}
	}

	while (!bStopping)
	{
		if (!OwnerComponent.IsValid() || !SphinxDecoder)
		{
			bStopping = true;
			continue;
		}

		// Dequeue a single chunk of audio. This will block if the queue is empty.
		TArray<int16> AudioChunk;
		if (OwnerComponent->GetAudioQueue().Dequeue(AudioChunk))
		{
			// Perform resampling if needed, directly on the chunk.
			int32 SampleRate = OwnerComponent->GetSourceSampleRate();
			if (SampleRate != PocketSphinxSampleRate)
			{
				SCOPE_CYCLE_COUNTER(STAT_InworldLipsync_DownsampleWaveform);
				OwnerComponent->DownsampleWaveform(AudioChunk, SampleRate, PocketSphinxSampleRate);
			}

			// Append the processed chunk to our local buffer
			ProcessingBuffer.Append(AudioChunk);

			// Define a threshold for how much audio to process at once.
			// Pocketsphinx works better on slightly larger chunks (e.g., 100-250ms).
			// 16000 (rate) * 0.1 (100ms) = 1600 samples.
			const int32 ProcessThreshold = 1600;

			if (ProcessingBuffer.Num() >= ProcessThreshold)
			{
				{
					SCOPE_CYCLE_COUNTER(STAT_InworldLipsync_SphinxProcess);
					ps_start_utt(SphinxDecoder);
					ps_process_raw(SphinxDecoder, ProcessingBuffer.GetData(), ProcessingBuffer.Num(), 0, 0);
					ps_end_utt(SphinxDecoder);
				}

				SCOPE_CYCLE_COUNTER(STAT_InworldLipsync_SphinxHyp);
				char const* hyp = ps_get_hyp(SphinxDecoder, nullptr);
				if (hyp)
				{
					// Convert result to FString
					FString PhonemeString(ANSI_TO_TCHAR(hyp));
					TArray<FString> Phonemes;
					PhonemeString.ParseIntoArray(Phonemes, TEXT(" "), true);

					if (Phonemes.Num() > 0)
					{
						FString LastPhoneme = Phonemes.Last();
						EViseme NewViseme = OwnerComponent->MapPhonemeToViseme(LastPhoneme);

						// Send the result back to the game thread
						AsyncTask(ENamedThreads::GameThread,
							[WeakOwner = OwnerComponent, NewViseme]()
							{
								if (WeakOwner.IsValid())
								{
									WeakOwner->UpdateVisemeFromWorker(NewViseme);
								}
							});

					}
				}

				ProcessingBuffer.Empty();

			}
		}
		else
		{
			if (bVoiceComplete)
			{
				bVoiceComplete = false;
				ResetBuffer();
				EViseme NewViseme = EViseme::VS_Silence;

				AsyncTask(ENamedThreads::GameThread,
					[WeakOwner = OwnerComponent, NewViseme]()
					{
						if (WeakOwner.IsValid())
						{
							WeakOwner->UpdateVisemeFromWorker(NewViseme);
						}
					});
			}
			// If the queue is empty, sleep for a moment to avoid a tight loop spinning the CPU.
			FPlatformProcess::Sleep(0.01f);	   // Sleep for 10ms
		}
	}

	return 0;
}

void FAudioProcessingRunnable::Stop()
{
	bStopping = true;
}

void FAudioProcessingRunnable::VoiceCompleted()
{
	bVoiceComplete = true;
}


void FAudioProcessingRunnable::Exit()
{
	// Clean up any resources here
	ShutdownPocketSphinx();
	UE_LOG(LogInworldLipsync, Log, TEXT("Audio Processing Thread Exiting."));
}

void FAudioProcessingRunnable::ResetBuffer()
{
	ProcessingBuffer.Empty();
}

void FAudioProcessingRunnable::EnsureCompletion()
{
	Stop();
	if (Thread)
	{
		// This blocks the calling thread (e.g., the game thread during EndPlay)
		// until our worker thread's Run() method has finished.
		Thread->WaitForCompletion();
	}
}

UInworldLipsyncComponent::UInworldLipsyncComponent()
{
	PrimaryComponentTick.bCanEverTick = false;
}

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

	// Bind to the SignedAudioData delegate.
	if (UInworldVoiceAudioComponent* VoiceComponent = GetOwner()->GetComponentByClass<UInworldVoiceAudioComponent>())
	{
		VoiceComponent->OnSignedAudioDataGenerated =
			FOnSignedAudioDataGenerated::CreateUObject(this, &UInworldLipsyncComponent::ProcessAudioFrame);

		AudioDelegateHandle =
			VoiceComponent->OnVoiceAudioCompleteNative.AddUObject(this, &UInworldLipsyncComponent::OnVoiceComplete);
	}

	// --- Create our phoneme-to-viseme mapping ---
	PhonemeToVisemeMap.Add(TEXT("P"), EViseme::VS_PP);
	PhonemeToVisemeMap.Add(TEXT("B"), EViseme::VS_PP);
	PhonemeToVisemeMap.Add(TEXT("M"), EViseme::VS_PP);
	PhonemeToVisemeMap.Add(TEXT("F"), EViseme::VS_FF);
	PhonemeToVisemeMap.Add(TEXT("V"), EViseme::VS_FF);
	PhonemeToVisemeMap.Add(TEXT("TH"), EViseme::VS_TH);
	PhonemeToVisemeMap.Add(TEXT("DH"), EViseme::VS_TH);
	PhonemeToVisemeMap.Add(TEXT("T"), EViseme::VS_DD);
	PhonemeToVisemeMap.Add(TEXT("D"), EViseme::VS_DD);
	PhonemeToVisemeMap.Add(TEXT("K"), EViseme::VS_kk);
	PhonemeToVisemeMap.Add(TEXT("G"), EViseme::VS_kk);
	PhonemeToVisemeMap.Add(TEXT("CH"), EViseme::VS_CH);
	PhonemeToVisemeMap.Add(TEXT("JH"), EViseme::VS_CH);
	PhonemeToVisemeMap.Add(TEXT("S"), EViseme::VS_SS);
	PhonemeToVisemeMap.Add(TEXT("Z"), EViseme::VS_SS);
	PhonemeToVisemeMap.Add(TEXT("N"), EViseme::VS_nn);
	PhonemeToVisemeMap.Add(TEXT("L"), EViseme::VS_nn);
	PhonemeToVisemeMap.Add(TEXT("R"), EViseme::VS_RR);
	PhonemeToVisemeMap.Add(TEXT("AA"), EViseme::VS_aa);
	PhonemeToVisemeMap.Add(TEXT("AE"), EViseme::VS_aa);
	PhonemeToVisemeMap.Add(TEXT("AH"), EViseme::VS_aa);
	PhonemeToVisemeMap.Add(TEXT("EH"), EViseme::VS_E);
	PhonemeToVisemeMap.Add(TEXT("ER"), EViseme::VS_E);
	PhonemeToVisemeMap.Add(TEXT("IH"), EViseme::VS_I);
	PhonemeToVisemeMap.Add(TEXT("IY"), EViseme::VS_I);
	PhonemeToVisemeMap.Add(TEXT("OW"), EViseme::VS_O);
	PhonemeToVisemeMap.Add(TEXT("UH"), EViseme::VS_U);
	PhonemeToVisemeMap.Add(TEXT("UW"), EViseme::VS_U);
	PhonemeToVisemeMap.Add(TEXT("SIL"), EViseme::VS_Silence);	 // Silence

	// InitializePocketSphinx();

	// Create and start the worker thread.
	AudioProcessor = MakeUnique<FAudioProcessingRunnable>(this);
}

void UInworldLipsyncComponent::EndPlay(const EEndPlayReason::Type EndPlayReason)
{
	if (AudioProcessor)
	{
		AudioProcessor->EnsureCompletion();
		AudioProcessor->ResetBuffer();
		AudioProcessor.Reset();	   // Explicitly destroy the runnable.
	}

	// Unbind the SignedAudioData delegate.
	if (UInworldVoiceAudioComponent* VoiceComponent = GetOwner()->GetComponentByClass<UInworldVoiceAudioComponent>())
	{
		VoiceComponent->OnVoiceAudioCompleteNative.Remove(AudioDelegateHandle);
		VoiceComponent->OnSignedAudioDataGenerated.Unbind();
	}
	ResetBuffer();
	// ShutdownPocketSphinx();
	Super::EndPlay(EndPlayReason);
}

FString VisemeToString(EViseme Viseme)
{
	FString VisemeString(TEXT("SIL"));
	switch (Viseme)
	{
		case EViseme::VS_PP:
			VisemeString = FString(TEXT("PP"));
			break;

		case EViseme::VS_FF:
			VisemeString = FString(TEXT("FF"));
			break;

		case EViseme::VS_TH:
			VisemeString = FString(TEXT("TH"));
			break;

		case EViseme::VS_DD:
			VisemeString = FString(TEXT("DD"));
			break;

		case EViseme::VS_kk:
			VisemeString = FString(TEXT("Kk"));
			break;

		case EViseme::VS_CH:
			VisemeString = FString(TEXT("CH"));
			break;

		case EViseme::VS_SS:
			VisemeString = FString(TEXT("SS"));
			break;

		case EViseme::VS_nn:
			VisemeString = FString(TEXT("nn"));
			break;

		case EViseme::VS_RR:
			VisemeString = FString(TEXT("RR"));
			break;

		case EViseme::VS_aa:
			VisemeString = FString(TEXT("Aa"));
			break;
		case EViseme::VS_E:
			VisemeString = FString(TEXT("E"));
			break;
		case EViseme::VS_I:
			VisemeString = FString(TEXT("I"));
			break;
		case EViseme::VS_O:
			VisemeString = FString(TEXT("O"));
			break;
		case EViseme::VS_U:
			VisemeString = FString(TEXT("U"));
			break;
	}
	return VisemeString;
}

void UInworldLipsyncComponent::OnVoiceComplete(UInworldVoiceAudioComponent* AudioComponent, const FInworldData_TTSOutput& Voice)
{
	if (AudioProcessor)
	{
		AudioProcessor->VoiceCompleted();
	}
}

void UInworldLipsyncComponent::ResetBuffer()
{
	if (AudioProcessor)
	{
		AudioProcessor->ResetBuffer();
	}
	CurrentViseme = EViseme::VS_Silence;
	AudioQueue.Empty();
}

// NOTE: Resample is very costly to run.
bool UInworldLipsyncComponent::ResampleWaveform(TArray<int16>& Waveform, int32 CurrentSampleRate, int32 TargetSampleRate)
{
	// --- Initial validation checks ---
	if (CurrentSampleRate <= 0 || TargetSampleRate <= 0 || Waveform.IsEmpty())
	{
		UE_LOG(LogInworldLipsync, Error, TEXT("DownsampleWaveform failed: SampleRates must be positive & waveform not empty."));
		return false;
	}

	// No resampling needed if rates are the same
	if (CurrentSampleRate == TargetSampleRate)
	{
		return true;
	}

	// --- Linear Interpolation Resampling ---

	const double Ratio = static_cast<double>(CurrentSampleRate) / static_cast<double>(TargetSampleRate);
	const int32 NewSize = FMath::Floor(Waveform.Num() / Ratio);

	// Create a new array to hold the resampled data.
	TArray<int16> ResampledWaveform;
	ResampledWaveform.Reserve(NewSize);	   // Pre-allocate memory for efficiency

	for (int32 i = 0; i < NewSize; ++i)
	{
		const double SourceIndex = i * Ratio;
		const int32 IndexFloor = static_cast<int32>(FMath::Floor(SourceIndex));

		// Ensure the ceiling index does not go out of bounds
		const int32 IndexCeil = FMath::Min(IndexFloor + 1, Waveform.Num() - 1);

		// Get the fractional part for interpolation (the 'alpha' in a Lerp)
		const float InterpFraction = static_cast<float>(SourceIndex - IndexFloor);

		// Get the two samples to interpolate between
		const float SampleA = static_cast<float>(Waveform[IndexFloor]);
		const float SampleB = static_cast<float>(Waveform[IndexCeil]);

		// Perform linear interpolation (Lerp)
		const float InterpolatedSampleFloat = FMath::Lerp(SampleA, SampleB, InterpFraction);

		// Clamp the result to the valid range of int16 to prevent overflow, then cast.
		const int16 FinalSample = static_cast<int16>(FMath::Clamp(InterpolatedSampleFloat, (float) INT16_MIN, (float) INT16_MAX));

		ResampledWaveform.Add(FinalSample);
	}

	// Replace the original waveform data with the new resampled data
	// MoveTemp is an efficient way to transfer ownership of the array's data
	Waveform = MoveTemp(ResampledWaveform);

	return true;
}

// default 24000hz in InworldTTS.
bool UInworldLipsyncComponent::Resample24kTo16k(TArray<int16>& Waveform)
{
	if (Waveform.IsEmpty())
	{
		return true;
	}

	// For every 3 input samples, we will generate 2 output samples.
	const int32 NewSize = (Waveform.Num() * 2) / 3;

	if (NewSize == 0)
	{
		Waveform.Empty();
		return true;
	}

	// The number of 2-sample output blocks we need to create.
	const int32 NumBlocks = NewSize / 2;

	for (int32 i = 0; i < NumBlocks; ++i)
	{
		const int32 WriteIndex = i * 2;
		const int32 ReadIndex = i * 3;

		// Output Sample 1: Take the first sample from the group of three directly.
		Waveform[WriteIndex] = Waveform[ReadIndex];

		// Output Sample 2: Average the next two samples.
		// We cast to int32 for the sum to prevent potential overflow before dividing.
		const int32 Sum = static_cast<int32>(Waveform[ReadIndex + 1]) + static_cast<int32>(Waveform[ReadIndex + 2]);
		Waveform[WriteIndex + 1] = static_cast<int16>(Sum / 2);
	}

	// Truncate the array to its new, smaller size without shrinking.
#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION < 4
	Waveform.SetNum(NewSize, false);
#else
	Waveform.SetNum(NewSize, EAllowShrinking::No);
#endif

	return true;
}

bool UInworldLipsyncComponent::DownsampleWaveform(TArray<int16>& Waveform, int32 CurrentSampleRate, int32 TargetSampleRate)
{
	if (CurrentSampleRate <= 0 || TargetSampleRate <= 0 || Waveform.IsEmpty())
	{
		UE_LOG(LogInworldLipsync, Error, TEXT("DownsampleWaveform failed: SampleRates must be positive & waveform not empty."));
		return false;
	}
	if (CurrentSampleRate == TargetSampleRate)
	{
		// No resampling needed if rates are the same
		return true;
	}
	if (CurrentSampleRate == 24000 && TargetSampleRate == 16000)
	{
		return Resample24kTo16k(Waveform);
	}

	// This method requires an integer downsampling ratio
	if (CurrentSampleRate <= TargetSampleRate || CurrentSampleRate % TargetSampleRate != 0)
	{
		UE_LOG(LogInworldLipsync, Warning,
			TEXT("In-place resampling requires the source sample rate to be an integer multiple of the target. Switching to costly "
				 "resample for %d Hz to %d Hz."),
			CurrentSampleRate, TargetSampleRate);
		return ResampleWaveform(Waveform, CurrentSampleRate, TargetSampleRate);
	}

	const int32 Ratio = CurrentSampleRate / TargetSampleRate;
	const int32 NewSize = Waveform.Num() / Ratio;

	if (NewSize == 0)
	{
		Waveform.Empty();
		return true;
	}

	for (int32 i = 0; i < NewSize; ++i)
	{
		const int32 TargetIndex = i * Ratio;

		if (TargetIndex > 0 && TargetIndex < Waveform.Num() - 1)
		{
			const int32 SmoothedSample =
				(static_cast<int32>(Waveform[TargetIndex - 1]) + static_cast<int32>(Waveform[TargetIndex]) +
					static_cast<int32>(Waveform[TargetIndex + 1])) /
				3;

			// Write the new, filtered sample to its final destination
			Waveform[i] = static_cast<int16>(SmoothedSample);
		}
		else
		{
			// Handle the edge case without filtering
			Waveform[i] = Waveform[TargetIndex];
		}
	}

#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION < 4
	Waveform.SetNum(NewSize, false);
#else
	Waveform.SetNum(NewSize, EAllowShrinking::No);
#endif


	return true;
}

void UInworldLipsyncComponent::ProcessAudioFrame(TArray<int16>&& AudioData, const int32 InSourceSampleRate)
{
	// Queue is thread safe and doesn't require a lock.
	{
		SCOPE_CYCLE_COUNTER(STAT_InworldLipsync_QueueAudio);
		AudioQueue.Enqueue(MoveTemp(AudioData));
		SourceSampleRate.store(InSourceSampleRate, std::memory_order_relaxed);
	}
}

EViseme UInworldLipsyncComponent::MapPhonemeToViseme(const FString& Phoneme)
{
	const EViseme* FoundViseme = PhonemeToVisemeMap.Find(Phoneme);
	return FoundViseme ? *FoundViseme : EViseme::VS_Silence;
}

void UInworldLipsyncComponent::UpdateVisemeFromWorker(EViseme NewViseme)
{
	CurrentViseme = NewViseme;

	// Since we are on the Game Thread, it's safe to update state and broadcast delegates.
	FInworldVisemeBlends VisemeBlend;
	VisemeBlend[VisemeToString(CurrentViseme)] = 1.f;
	OnVisemeUpdate.Broadcast(VisemeBlend);
}
