// Georgy Treshchev 2026 All Rights Reserved.

#include "RuntimeVoskSpeechRecognizerProvider.h"

#include "RuntimeSpeechRecognizerVosk.h"
#include "RuntimeVoskModelLibrary.h"
#include "RuntimeVoskRecognizerThread.h"
#include "SpeechRecognizer.h"

#if PLATFORM_ANDROID
#include "Async/Async.h"
#endif

URuntimeVoskSpeechRecognizerProvider::URuntimeVoskSpeechRecognizerProvider()
{
	if (!HasAnyFlags(RF_ClassDefaultObject))
	{
		Initialize();
	}
}

void URuntimeVoskSpeechRecognizerProvider::BeginDestroy()
{
	if (Thread.IsValid())
	{
		Thread->StopThread();
		Thread.Reset();
	}
	Super::BeginDestroy();
}

void URuntimeVoskSpeechRecognizerProvider::Initialize()
{
	Thread = MakeShared<FRuntimeVoskRecognizerThread>();
	ensure(Thread.IsValid());

	if (!Thread.IsValid())
	{
		return;
	}

	Thread->OnRecognitionFinished.AddWeakLambda(this, [this]()
	{
		OnRecognitionFinished.Broadcast();
	});

	Thread->OnRecognitionError.AddWeakLambda(this, [this](const FString& ShortErrorMessage, const FString& LongErrorMessage)
	{
		OnRecognitionError.Broadcast(ShortErrorMessage, LongErrorMessage);
	});

	Thread->OnRecognizedTextSegment.AddWeakLambda(this, [this](const FString& RecognizedWords)
	{
		OnRecognizedTextSegment.Broadcast(RecognizedWords);
	});

	Thread->OnRecognitionProgress.AddWeakLambda(this, [this](int32 Progress)
	{
		OnRecognitionProgress.Broadcast(Progress);
	});

	Thread->OnRecognitionStopped.AddWeakLambda(this, [this]()
	{
		OnRecognitionStopped.Broadcast();
	});
}

bool URuntimeVoskSpeechRecognizerProvider::SetVoskModelPath(ESpeechRecognizerLanguage Language, const FString& ModelFolderPath)
{
	if (ModelFolderPath.IsEmpty())
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("Unable to set the Vosk model path: Provided path is empty"));
		return false;
	}

	if (!Thread.IsValid() || !Thread->GetIsStopped())
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("Unable to set the Vosk model path while recognition is running. Call StopSpeechRecognition() first"));
		return false;
	}

	VoskModelPathsByLanguage.Add(Language, ModelFolderPath);

	if (RecognitionParameters.Language == Language)
	{
		Thread->SetModelFolderPath(ModelFolderPath);
	}

	UE_LOG(LogRuntimeSpeechRecognizerVosk, Log, TEXT("Successfully set Vosk model path for language '%s': %s"),
		*USpeechRecognizer::GetLanguageCode(Language), *ModelFolderPath);
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetVoskModelPathByName(ESpeechRecognizerLanguage Language, const FName ModelName)
{
	if (ModelName.IsNone())
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("Unable to set the Vosk model path: No model name specified"));
		return false;
	}
 
	const FString ModelId = ModelName.ToString();
 
#if PLATFORM_ANDROID
	if (!URuntimeVoskModelLibrary::IsVoskModelOnDisk(ModelId))
	{
		if (!URuntimeVoskModelLibrary::IsVoskModelInPackagedContent(ModelId))
		{
			UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("Unable to set the Vosk model path: Model '%s' is not present in the packaged build"), *ModelId);
			return false;
		}

		URuntimeVoskModelLibrary::MakeVoskModelReadyOnAndroid(ModelId);
		PendingAndroidModelCopies.Add(ModelId);
 
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Log, TEXT("Vosk model '%s' is being prepared in the background for language '%s'; "
			"StartSpeechRecognition() will wait for it automatically if it isn't ready yet"),
			*ModelId, *USpeechRecognizer::GetLanguageCode(Language));
	}
#endif
 
	return SetVoskModelPath(Language, URuntimeVoskModelLibrary::GetVoskModelFolderPath(ModelId));
}

void URuntimeVoskSpeechRecognizerProvider::StartSpeechRecognition(const FOnSpeechRecognitionStartedStatic& OnStarted)
{
	if (!Thread.IsValid())
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("Failed to start speech recognition: Vosk provider thread is not valid"));
		OnStarted.ExecuteIfBound(false);
		return;
	}
 
	const FString* ModelPath = VoskModelPathsByLanguage.Find(RecognitionParameters.Language);
	if (!ModelPath || ModelPath->IsEmpty())
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("Failed to start speech recognition: No Vosk model path set for language '%s'. Call SetVoskModelPath() first"),
			*USpeechRecognizer::GetLanguageCode(RecognitionParameters.Language));
		OnStarted.ExecuteIfBound(false);
		return;
	}
 
#if PLATFORM_ANDROID
	const FString ModelId = FPaths::GetCleanFilename(*ModelPath);
	if (PendingAndroidModelCopies.Contains(ModelId))
	{
		TWeakObjectPtr<URuntimeVoskSpeechRecognizerProvider> WeakThis(this);
		const FString CapturedModelPath = *ModelPath;
 
		URuntimeVoskModelLibrary::MakeVoskModelReadyOnAndroid(ModelId).Next(
			[WeakThis, OnStarted, ModelId, CapturedModelPath](bool bReady)
		{
			URuntimeVoskSpeechRecognizerProvider* StrongThis = WeakThis.Get();
			if (!StrongThis)
			{
				OnStarted.ExecuteIfBound(false);
				return;
			}
 
			StrongThis->PendingAndroidModelCopies.Remove(ModelId);
 
			if (!bReady)
			{
				UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("Failed to start speech recognition: Vosk model '%s' could not be prepared on this device"), *ModelId);
				OnStarted.ExecuteIfBound(false);
				return;
			}

			if (!StrongThis->Thread.IsValid())
			{
				OnStarted.ExecuteIfBound(false);
				return;
			}
 
			StrongThis->Thread->SetModelFolderPath(CapturedModelPath);
			StrongThis->Thread->SetStepSizeMs(StrongThis->RecognitionParameters.StepSizeMs);
			StrongThis->Thread->SetMaxTokensPerSegment(StrongThis->RecognitionParameters.MaxTokens);
			StrongThis->Thread->StartThread().Next([OnStarted](bool bSucceeded) { OnStarted.ExecuteIfBound(bSucceeded); });
		});
		return;
	}
#endif
 
	Thread->SetModelFolderPath(*ModelPath);
	Thread->SetStepSizeMs(RecognitionParameters.StepSizeMs);
	Thread->SetMaxTokensPerSegment(RecognitionParameters.MaxTokens);
 
	Thread->StartThread().Next([OnStarted](bool bSucceeded) { OnStarted.ExecuteIfBound(bSucceeded); });
}

void URuntimeVoskSpeechRecognizerProvider::StopSpeechRecognition()
{
	if (Thread.IsValid())
	{
		Thread->StopThread();
	}
}

void URuntimeVoskSpeechRecognizerProvider::ProcessAudioData(Audio::FAlignedFloatBuffer PCMData, float SampleRate, int32 NumOfChannels, bool bLast)
{
	if (Thread.IsValid())
	{
		Thread->ProcessPCMData(MoveTemp(PCMData), SampleRate, NumOfChannels, bLast);
	}
}

void URuntimeVoskSpeechRecognizerProvider::ForceProcessPendingAudioData()
{
	if (Thread.IsValid())
	{
		Thread->ForceProcessPendingAudioData();
	}
}

void URuntimeVoskSpeechRecognizerProvider::ClearAudioData(bool bClearPendingAudioData, bool bClearAudioQueue)
{
	if (Thread.IsValid())
	{
		Thread->ClearAudioData(bClearPendingAudioData, bClearAudioQueue);
	}
}

bool URuntimeVoskSpeechRecognizerProvider::GetIsStopped() const
{
	return Thread.IsValid() ? Thread->GetIsStopped() : true;
}

bool URuntimeVoskSpeechRecognizerProvider::GetIsStopping() const
{
	return Thread.IsValid() ? Thread->GetIsStopping() : false;
}

bool URuntimeVoskSpeechRecognizerProvider::GetIsFinished() const
{
	return Thread.IsValid() ? Thread->GetIsFinished() : true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetRecognitionParameters(const FSpeechRecognitionParameters& Parameters)
{
	if (!GetIsStopped())
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Error, TEXT("Unable to set recognition parameters while the Vosk recognizer is running"));
		return false;
	}

	RecognitionParameters = Parameters;

	if (Thread.IsValid())
	{
		Thread->SetStepSizeMs(RecognitionParameters.StepSizeMs);
		Thread->SetMaxTokensPerSegment(RecognitionParameters.MaxTokens);
	}

	// Reflect unsupported fields being ignored, at Verbose level, without failing the call
	if (Parameters.bTranslateToEnglish)
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("bTranslateToEnglish is not supported by the Vosk provider and will be ignored"));
	}
	if (Parameters.BeamSize > 0)
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("BeamSize is not supported by the Vosk provider and will be ignored"));
	}

	return true;
}

FSpeechRecognitionParameters URuntimeVoskSpeechRecognizerProvider::GetRecognitionParameters() const
{
	return RecognitionParameters;
}

ESpeechRecognizerLanguage URuntimeVoskSpeechRecognizerProvider::GetDetectedLanguage() const
{
	// Vosk selects a model per-language rather than detecting it at runtime, so the "detected" language
	// is simply whichever language's model is currently configured
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("Vosk does not perform runtime language detection; returning the currently configured recognition language"));
	return RecognitionParameters.Language;
}

bool URuntimeVoskSpeechRecognizerProvider::SetLanguage(ESpeechRecognizerLanguage Language)
{
	if (Language == ESpeechRecognizerLanguage::Auto)
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Warning, TEXT("Vosk does not support automatic language detection ('Auto'). Ignoring and keeping the current language"));
		return true;
	}

	RecognitionParameters.Language = Language;

	const FString* ModelPath = VoskModelPathsByLanguage.Find(Language);
	if (ModelPath && Thread.IsValid())
	{
		Thread->SetModelFolderPath(*ModelPath);
	}
	else
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Warning, TEXT("SetLanguage('%s') called, but no Vosk model path has been set for this language yet. Call SetVoskModelPath() before starting recognition"),
			*USpeechRecognizer::GetLanguageCode(Language));
	}

	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetStepSize(int32 Value)
{
	RecognitionParameters.StepSizeMs = Value;
	if (Thread.IsValid())
	{
		Thread->SetStepSizeMs(Value);
	}
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetNonStreamingDefaults()
{
	RecognitionParameters.SetNonStreamingDefaults();
	SetMaxTokens(RecognitionParameters.MaxTokens);
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetNonStreamingDefaults() has limited effect on the Vosk provider, which is inherently a streaming recognizer"));
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetStreamingDefaults()
{
	RecognitionParameters.SetStreamingDefaults();
	SetStepSize(200);
	SetMaxTokens(RecognitionParameters.MaxTokens);
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetNumOfThreads(int32 Value)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetNumOfThreads(%d) is not supported by the Vosk provider and will be ignored"), Value);
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetTranslateToEnglish(bool bTranslate)
{
	if (bTranslate)
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Warning, TEXT("SetTranslateToEnglish(true) is not supported by the Vosk provider and will be ignored"));
	}
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetNoContext(bool bNoContext)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetNoContext(%s) is not applicable to the Vosk provider and will be ignored"), bNoContext ? TEXT("true") : TEXT("false"));
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetSingleSegment(bool bSingleSegment)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetSingleSegment(%s) is not applicable to the Vosk provider and will be ignored"), bSingleSegment ? TEXT("true") : TEXT("false"));
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetMaxTokens(int32 Value)
{
	RecognitionParameters.MaxTokens = Value;
	if (Thread.IsValid())
	{
		Thread->SetMaxTokensPerSegment(Value);
	}
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetMaxTokens(%d): Vosk has no native token limit; this is used to select streaming (>0) vs accumulated (0) text delivery"), Value);
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetSpeedUp(bool bSpeedUp)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetSpeedUp(%s) is not supported by the Vosk provider and will be ignored"), bSpeedUp ? TEXT("true") : TEXT("false"));
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetAudioContextSize(int32 Value)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetAudioContextSize(%d) is not supported by the Vosk provider and will be ignored"), Value);
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetTemperatureToIncrease(float Value)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetTemperatureToIncrease(%f) is not supported by the Vosk provider and will be ignored"), Value);
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetEntropyThreshold(float Value)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetEntropyThreshold(%f) is not supported by the Vosk provider and will be ignored"), Value);
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetSuppressBlank(bool Value)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetSuppressBlank(%s) is not supported by the Vosk provider and will be ignored"), Value ? TEXT("true") : TEXT("false"));
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetSuppressNonSpeechTokens(bool Value)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetSuppressNonSpeechTokens(%s) is not supported by the Vosk provider and will be ignored"), Value ? TEXT("true") : TEXT("false"));
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetBeamSize(int32 Value)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetBeamSize(%d) is not supported by the Vosk provider and will be ignored"), Value);
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetInitialPrompt(const FString& Value)
{
	if (!Value.IsEmpty())
	{
		UE_LOG(LogRuntimeSpeechRecognizerVosk, Warning, TEXT("SetInitialPrompt() is not supported by the Vosk provider and will be ignored"));
	}
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetGpuAcceleration(bool Value)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetGpuAcceleration(%s) is not applicable to the Vosk provider (CPU-only) and will be ignored"), Value ? TEXT("true") : TEXT("false"));
	return true;
}

bool URuntimeVoskSpeechRecognizerProvider::SetGpuDeviceId(int32 Value)
{
	UE_LOG(LogRuntimeSpeechRecognizerVosk, Verbose, TEXT("SetGpuDeviceId(%d) is not applicable to the Vosk provider (CPU-only) and will be ignored"), Value);
	return true;
}