// Georgy Treshchev 2026 All Rights Reserved.

#include "RuntimeVoskRecognizerThread.h"
#include "RuntimeSpeechRecognizerVosk.h"
#include "AudioResampler.h"
#include "SampleBuffer.h"
#include "HAL/RunnableThread.h"
#include "Async/Async.h"
#include "Serialization/JsonSerializer.h"
#include "Dom/JsonObject.h"

#include "vosk_api.h"

#if PLATFORM_ANDROID
#include "Android/AndroidJNI.h"
#include "Android/AndroidApplication.h"
#include "Runtime/Launch/Resources/Version.h"
 
FString RuntimeVoskAndroidNoBackupFilesDir;
bool RuntimeVoskAndroidNoBackupFilesDirValid = false;
 
#if USE_ANDROID_JNI
#if ENGINE_MAJOR_VERSION >= 5
JNI_METHOD void Java_com_epicgames_unreal_GameActivity_nativeRuntimeVoskSetNoBackupFilesDir(JNIEnv* jenv, jobject thiz, jboolean isValid, jstring noBackupFilesDir)
#else
JNI_METHOD void Java_com_epicgames_ue4_GameActivity_nativeRuntimeVoskSetNoBackupFilesDir(JNIEnv* jenv, jobject thiz, jboolean isValid, jstring noBackupFilesDir)
#endif
{
	RuntimeVoskAndroidNoBackupFilesDir = FJavaHelper::FStringFromParam(jenv, noBackupFilesDir);
	RuntimeVoskAndroidNoBackupFilesDirValid = isValid;
}
#endif
#endif

namespace RuntimeVosk
{
	constexpr int32 VoskSampleRate = 16000;
	constexpr int32 VoskFeedChunkSize = 1600;

	static bool ExtractText(const FString& JsonString, FString& OutText, bool& bOutIsPartial)
	{
		TSharedPtr<FJsonObject> JsonObject;
		const TSharedRef<TJsonReader<>> JsonReader = TJsonReaderFactory<>::Create(JsonString);

		if (!FJsonSerializer::Deserialize(JsonReader, JsonObject) || !JsonObject.IsValid())
		{
			return false;
		}

		if (JsonObject->TryGetStringField(TEXT("text"), OutText))
		{
			bOutIsPartial = false;
			return OutText.Len() > 0;
		}

		if (JsonObject->TryGetStringField(TEXT("partial"), OutText))
		{
			bOutIsPartial = true;
			return OutText.Len() > 0;
		}

		return false;
	}
}

FRuntimeVoskRecognizerThread::FRuntimeVoskRecognizerThread()
	: bIsStopped(true)
	, bIsFinished(true)
	, bIsStopping(false)
	, VoskModelPtr(nullptr)
	, VoskRecognizerPtr(nullptr)
{
}

FRuntimeVoskRecognizerThread::~FRuntimeVoskRecognizerThread()
{
	ReleaseMemory();
}

void FRuntimeVoskRecognizerThread::SetModelFolderPath(const FString& InModelFolderPath)
{
	ModelFolderPath = InModelFolderPath;
}

void FRuntimeVoskRecognizerThread::SetStepSizeMs(int32 Value)
{
	StepSizeMs = Value;
}

void FRuntimeVoskRecognizerThread::SetMaxTokensPerSegment(int32 Value)
{
	MaxTokensPerSegment = Value;
}

TFuture<bool> FRuntimeVoskRecognizerThread::StartThread()
{
	if (!bIsStopped)
	{
		ReportError(TEXT("Thread start failed"), TEXT("Unable to start an already running Vosk recognizer thread"));
		return MakeFulfilledPromise<bool>(false).GetFuture();
	}

	if (bIsStopping)
	{
		ReportError(TEXT("Thread start failed"), TEXT("Unable to start a Vosk recognizer thread that is stopping"));
		return MakeFulfilledPromise<bool>(false).GetFuture();
	}

	if (ModelFolderPath.IsEmpty())
	{
		ReportError(TEXT("Vosk recognizer initialization failed"), TEXT("No Vosk model path was set. Call SetVoskModelPath() on the provider before starting recognition"));
		return MakeFulfilledPromise<bool>(false).GetFuture();
	}

	if (StartThreadPromise.IsValid())
	{
		ReportError(TEXT("Thread start failed"), TEXT("Unable to start a Vosk recognizer thread that is already starting"));
		return MakeFulfilledPromise<bool>(false).GetFuture();
	}

	StartThreadPromise = MakeUnique<TPromise<bool>>();
	TFuture<bool> ResultFuture = StartThreadPromise->GetFuture();

	TSharedRef<FRuntimeVoskRecognizerThread> ThisShared = AsShared();

	AsyncTask(ENamedThreads::AnyBackgroundHiPriTask, [ThisShared]() mutable
	{
		VoskModel* LoadedModel = vosk_model_new(TCHAR_TO_ANSI(*ThisShared->ModelFolderPath));
		if (!LoadedModel)
		{
			ThisShared->ReportError(TEXT("Vosk model loading failed"), FString::Printf(TEXT("Failed to load the Vosk model from '%s'"), *ThisShared->ModelFolderPath));
			AsyncTask(ENamedThreads::GameThread, [ThisShared]()
			{
				ThisShared->StartThreadPromise->SetValue(false);
				ThisShared->StartThreadPromise.Reset();
			});
			return;
		}

		VoskRecognizer* NewRecognizer = vosk_recognizer_new(LoadedModel, RuntimeVosk::VoskSampleRate);
		if (!NewRecognizer)
		{
			vosk_model_free(LoadedModel);
			ThisShared->ReportError(TEXT("Vosk recognizer creation failed"), TEXT("Failed to create the Vosk recognizer instance"));
			AsyncTask(ENamedThreads::GameThread, [ThisShared]()
			{
				ThisShared->StartThreadPromise->SetValue(false);
				ThisShared->StartThreadPromise.Reset();
			});
			return;
		}

		AsyncTask(ENamedThreads::GameThread, [ThisShared, LoadedModel, NewRecognizer]() mutable
		{
			ThisShared->VoskModelPtr = LoadedModel;
			ThisShared->VoskRecognizerPtr = NewRecognizer;

			ThisShared->bIsStopping.AtomicSet(false);
			ThisShared->bIsStopped.AtomicSet(false);
			ThisShared->bIsFinished.AtomicSet(true);

			FRunnableThread* ThreadPtr = FRunnableThread::Create(ThisShared.operator->(), TEXT("RuntimeVoskRecognizerThread"), 0, TPri_Highest);
			if (!ThreadPtr)
			{
				ThisShared->ReportError(TEXT("Thread creation failed"), TEXT("Failed to create the thread for the Vosk speech recognizer"));
				ThisShared->StartThreadPromise->SetValue(false);
				ThisShared->StartThreadPromise.Reset();
				return;
			}

			ThisShared->Thread.Reset(ThreadPtr);
			ThisShared->StartThreadPromise->SetValue(true);
			ThisShared->StartThreadPromise.Reset();
		});
	});

	return ResultFuture;
}

void FRuntimeVoskRecognizerThread::StopThread()
{
	TSharedRef<FRuntimeVoskRecognizerThread> ThisShared = AsShared();
	AsyncTask(ENamedThreads::AnyBackgroundHiPriTask, [ThisShared]() mutable
	{
		ThisShared->ReleaseMemory();
	});
}

void FRuntimeVoskRecognizerThread::ProcessPCMData(Audio::FAlignedFloatBuffer PCMData, float SampleRate, uint32 NumOfChannels, bool bLast)
{
	if (bIsStopped)
	{
		ReportError(TEXT("Audio processing failed"), TEXT("The audio data could not be processed since the Vosk recognizer thread is stopped"));
		return;
	}
	if (bIsStopping)
	{
		ReportError(TEXT("Audio processing failed"), TEXT("The audio data could not be processed since the Vosk recognizer thread is stopping"));
		return;
	}
	if (SampleRate <= 0.0f || NumOfChannels == 0)
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("Invalid sample rate (%f) or number of channels (%d)"), SampleRate, NumOfChannels);
		return;
	}

	// Resample to 16kHz mono, which Vosk requires
	Audio::FAlignedFloatBuffer ProcessedPCMData = MoveTemp(PCMData);

	if (NumOfChannels > 1)
	{
		Audio::TSampleBuffer<float> SampleBufferMix(ProcessedPCMData, NumOfChannels, SampleRate);
		SampleBufferMix.MixBufferToChannels(1);
		ProcessedPCMData = Audio::FAlignedFloatBuffer(SampleBufferMix.GetData(), SampleBufferMix.GetNumSamples());
	}

	if (static_cast<int32>(SampleRate) != RuntimeVosk::VoskSampleRate)
	{
		Audio::FAlignedFloatBuffer ResampledPCMData;
		const Audio::FResamplingParameters ResampleParameters = {
			Audio::EResamplingMethod::Linear,
			1,
			SampleRate,
			static_cast<float>(RuntimeVosk::VoskSampleRate),
			ProcessedPCMData
		};
		ResampledPCMData.AddUninitialized(Audio::GetOutputBufferSize(ResampleParameters));
		Audio::FResamplerResults ResampleResults;
		ResampleResults.OutBuffer = &ResampledPCMData;

		if (!Audio::Resample(ResampleParameters, ResampleResults))
		{
			ReportError(TEXT("Audio processing failed"), TEXT("Unable to resample audio data to 16kHz for the Vosk recognizer"));
			return;
		}
		ProcessedPCMData = MoveTemp(ResampledPCMData);
	}

	FScopeLock Lock(&PendingAudioGuard);
	PendingAudioBuffer.Append(ProcessedPCMData);

	const int32 NumOfSamplesPerStep = StepSizeMs > 0 ? static_cast<int32>((1e-3 * StepSizeMs) * RuntimeVosk::VoskSampleRate) : 0;

	if (bLast || NumOfSamplesPerStep <= 0 || PendingAudioBuffer.Num() >= NumOfSamplesPerStep)
	{
		FQueuedAudioChunk Chunk;
		Chunk.Buffer = MoveTemp(PendingAudioBuffer);
		Chunk.bLast = bLast;
		AudioQueue.Enqueue(MoveTemp(Chunk));
		PendingAudioBuffer.Reset();
	}
}

void FRuntimeVoskRecognizerThread::ForceProcessPendingAudioData()
{
	FScopeLock Lock(&PendingAudioGuard);

	FQueuedAudioChunk Chunk;
	Chunk.Buffer = MoveTemp(PendingAudioBuffer);
	Chunk.bLast = true;
	AudioQueue.Enqueue(MoveTemp(Chunk));
	PendingAudioBuffer.Reset();
}

void FRuntimeVoskRecognizerThread::ClearAudioData(bool bClearPendingAudioData, bool bClearAudioQueue)
{
	if (bClearPendingAudioData)
	{
		FScopeLock Lock(&PendingAudioGuard);
		PendingAudioBuffer.Reset();
	}
	if (bClearAudioQueue)
	{
		AudioQueue.Empty();
	}
}

bool FRuntimeVoskRecognizerThread::Init()
{
	return FRunnable::Init();
}

uint32 FRuntimeVoskRecognizerThread::Run()
{
	while (!bIsStopped && !bIsStopping)
	{
		FQueuedAudioChunk DequeuedChunk;
		while (AudioQueue.Dequeue(DequeuedChunk))
		{
			Audio::FAlignedFloatBuffer QueuedBuffer = MoveTemp(DequeuedChunk.Buffer);

			if (QueuedBuffer.Num() == 0 && !DequeuedChunk.bLast)
			{
				continue;
			}

			bIsFinished.AtomicSet(false);

			if (VoskRecognizerPtr)
			{
				int32 Offset = 0;
				while (Offset < QueuedBuffer.Num())
				{
					const int32 CurrentChunkSize = FMath::Min(RuntimeVosk::VoskFeedChunkSize, QueuedBuffer.Num() - Offset);

					Audio::FAlignedFloatBuffer ScaledChunk;
					ScaledChunk.SetNumUninitialized(CurrentChunkSize);
					for (int32 i = 0; i < CurrentChunkSize; ++i)
					{
						ScaledChunk[i] = QueuedBuffer[Offset + i] * 32768.0f;
					}

					const int32 AcceptResult = vosk_recognizer_accept_waveform_f(VoskRecognizerPtr, ScaledChunk.GetData(), CurrentChunkSize);
					Offset += CurrentChunkSize;

					FString JsonResult;
					bool bIsPartial = false;
					if (AcceptResult > 0)
					{
						JsonResult = FString(UTF8_TO_TCHAR(vosk_recognizer_result(VoskRecognizerPtr)));
						bIsPartial = false;
						vosk_recognizer_reset(VoskRecognizerPtr);
					}
					else if (AcceptResult == 0)
					{
						JsonResult = FString(UTF8_TO_TCHAR(vosk_recognizer_partial_result(VoskRecognizerPtr)));
						bIsPartial = true;
					}
					else
					{
						ReportError(TEXT("Vosk recognition failed"), TEXT("vosk_recognizer_accept_waveform_f returned an error"));
						continue;
					}

					HandleVoskResultJson(JsonResult, bIsPartial);
				}
			}

			if (DequeuedChunk.bLast)
			{
				if (VoskRecognizerPtr)
				{
					const FString FinalJsonResult = FString(UTF8_TO_TCHAR(vosk_recognizer_final_result(VoskRecognizerPtr)));
					HandleVoskResultJson(FinalJsonResult, /*bIsPartial=*/false);
					vosk_recognizer_reset(VoskRecognizerPtr);
				}

				if (MaxTokensPerSegment <= 0)
				{
					FlushAccumulatedText();
				}

				bool bQueueIsEmpty = false;
				bool bPendingIsEmpty = false;
				{
					FScopeLock Lock(&PendingAudioGuard);
					bPendingIsEmpty = (PendingAudioBuffer.Num() == 0);
				}
				bQueueIsEmpty = AudioQueue.IsEmpty();

				if (bPendingIsEmpty && bQueueIsEmpty && !bIsFinished)
				{
					bIsFinished.AtomicSet(true);
					TSharedRef<FRuntimeVoskRecognizerThread> ThisShared = AsShared();
					AsyncTask(ENamedThreads::GameThread, [ThisShared]()
					{
						ThisShared->OnRecognitionProgress.Broadcast(100);
						ThisShared->OnRecognitionFinished.Broadcast();
					});
				}
			}
		}

		FPlatformProcess::Sleep(0.01f);
	}

	if (bIsStopping)
	{
		TSharedRef<FRuntimeVoskRecognizerThread> ThisShared = AsShared();
		AsyncTask(ENamedThreads::GameThread, [ThisShared]()
		{
			ThisShared->OnRecognitionStopped.Broadcast();
		});
	}

	return 0;
}

void FRuntimeVoskRecognizerThread::Stop()
{
	bIsStopping.AtomicSet(true);
	bIsFinished.AtomicSet(true);
	FRunnable::Stop();
}

void FRuntimeVoskRecognizerThread::Exit()
{
	bIsStopping.AtomicSet(false);
	bIsStopped.AtomicSet(true);
	bIsFinished.AtomicSet(true);
	FRunnable::Exit();
}

bool FRuntimeVoskRecognizerThread::GetIsStopped() const
{
	return bIsStopped;
}

bool FRuntimeVoskRecognizerThread::GetIsStopping() const
{
	return bIsStopping;
}

bool FRuntimeVoskRecognizerThread::GetIsFinished() const
{
	return bIsFinished;
}

void FRuntimeVoskRecognizerThread::HandleVoskResultJson(const FString& JsonResult, bool bIsPartial)
{
	FString RecognizedText;
	bool bActuallyPartial = bIsPartial;
	if (!RuntimeVosk::ExtractText(JsonResult, RecognizedText, bActuallyPartial))
	{
		return;
	}

	if (MaxTokensPerSegment > 0)
	{
		TSharedRef<FRuntimeVoskRecognizerThread> ThisShared = AsShared();
		AsyncTask(ENamedThreads::GameThread, [ThisShared, RecognizedText]()
		{
			UE_LOG(LogRuntimeSpeechRecognizerVosk, Log, TEXT("Recognized text segment: \"%s\""), *RecognizedText);
			ThisShared->OnRecognizedTextSegment.Broadcast(RecognizedText);
		});
		return;
	}

	if (!bActuallyPartial)
	{
		if (AccumulatedText.IsEmpty())
		{
			AccumulatedText = RecognizedText;
		}
		else
		{
			AccumulatedText += TEXT(" ");
			AccumulatedText += RecognizedText;
		}
	}
}

void FRuntimeVoskRecognizerThread::FlushAccumulatedText()
{
	if (AccumulatedText.IsEmpty())
	{
		return;
	}

	FString TextToSend = MoveTemp(AccumulatedText);
	AccumulatedText.Reset();

	TSharedRef<FRuntimeVoskRecognizerThread> ThisShared = AsShared();
	AsyncTask(ENamedThreads::GameThread, [ThisShared, TextToSend]()
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Log, TEXT("Recognized text segment (accumulated): \"%s\""), *TextToSend);
		ThisShared->OnRecognizedTextSegment.Broadcast(TextToSend);
	});
}

void FRuntimeVoskRecognizerThread::ReleaseMemory()
{
	Thread.Reset();

	if (VoskRecognizerPtr)
	{
		vosk_recognizer_free(VoskRecognizerPtr);
		VoskRecognizerPtr = nullptr;
	}
	if (VoskModelPtr)
	{
		vosk_model_free(VoskModelPtr);
		VoskModelPtr = nullptr;
	}

	AccumulatedText.Reset();

	bIsStopping.AtomicSet(false);
	bIsStopped.AtomicSet(true);
	bIsFinished.AtomicSet(true);
}

void FRuntimeVoskRecognizerThread::ReportError(const FString& ShortErrorMessage, const FString& LongErrorMessage)
{
	TSharedRef<FRuntimeVoskRecognizerThread> ThisShared = AsShared();
	AsyncTask(ENamedThreads::GameThread, [ThisShared, ShortErrorMessage, LongErrorMessage]()
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("%s: %s"), *ShortErrorMessage, *LongErrorMessage);
		ThisShared->OnRecognitionError.Broadcast(ShortErrorMessage, LongErrorMessage);
	});
}