// Georgy Treshchev 2026 All Rights Reserved.

#pragma once

#include "CoreMinimal.h"
#include "RuntimeSpeechRecognizerProviderBase.h"
#include "RuntimeVoskRecognizerThread.h"
#include "HAL/Runnable.h"
#include "HAL/ThreadSafeBool.h"
#include "Containers/Queue.h"
#include "SampleBuffer.h"
#include "Async/Future.h"

#include "RuntimeVoskSpeechRecognizerProvider.generated.h"

/**
 * Vosk API speech recognizer provider
 *
 * Ships in a separate extension plugin (RuntimeSpeechRecognizerVosk) which depends on the main
 * Runtime Speech Recognizer plugin, but is not depended on by it. Enable it at runtime by calling
 * USpeechRecognizer::SetSpeechRecognizerProvider(URuntimeVoskSpeechRecognizerProvider::StaticClass())
 *
 * The public behavior of USpeechRecognizer is unaffected by which provider is active - all setters
 * that Vosk does not support (e.g. GPU acceleration, beam size, entropy threshold) simply return true
 * and log a message, so existing Blueprint/C++ logic written against the whisper.cpp provider keeps
 * working unmodified after switching to Vosk
 */
UCLASS(BlueprintType, Category = "Runtime Speech Recognizer")
class RUNTIMESPEECHRECOGNIZERVOSK_API URuntimeVoskSpeechRecognizerProvider : public URuntimeSpeechRecognizerProviderBase
{
	GENERATED_BODY()

public:
	URuntimeVoskSpeechRecognizerProvider();

	//~ Begin UObject Interface
	virtual void BeginDestroy() override;
	//~ End UObject Interface

	//~ Begin URuntimeSpeechRecognizerProviderBase Interface
	virtual void StartSpeechRecognition(const FOnSpeechRecognitionStartedStatic& OnStarted) override;
	virtual void StopSpeechRecognition() override;
	virtual void ProcessAudioData(Audio::FAlignedFloatBuffer PCMData, float SampleRate, int32 NumOfChannels, bool bLast) override;
	virtual void ForceProcessPendingAudioData() override;
	virtual void ClearAudioData(bool bClearPendingAudioData, bool bClearAudioQueue) override;
	virtual bool GetIsStopped() const override;
	virtual bool GetIsStopping() const override;
	virtual bool GetIsFinished() const override;
	virtual bool SetRecognitionParameters(const FSpeechRecognitionParameters& Parameters) override;
	virtual FSpeechRecognitionParameters GetRecognitionParameters() const override;
	virtual ESpeechRecognizerLanguage GetDetectedLanguage() const override;
	virtual FString GetProviderName() const override { return TEXT("Vosk API"); }

	//~ Supported by Vosk
	virtual bool SetLanguage(ESpeechRecognizerLanguage Language) override;
	virtual bool SetStepSize(int32 Value) override;

	//~ Not natively supported by Vosk: return true (no-op success) + log, so switching providers
	//~ doesn't require the caller to change any logic. See class comment above
	virtual bool SetNonStreamingDefaults() override;
	virtual bool SetStreamingDefaults() override;
	virtual bool SetNumOfThreads(int32 Value) override;
	virtual bool SetTranslateToEnglish(bool bTranslate) override;
	virtual bool SetNoContext(bool bNoContext) override;
	virtual bool SetSingleSegment(bool bSingleSegment) override;
	virtual bool SetMaxTokens(int32 Value) override;
	virtual bool SetSpeedUp(bool bSpeedUp) override;
	virtual bool SetAudioContextSize(int32 Value) override;
	virtual bool SetTemperatureToIncrease(float Value) override;
	virtual bool SetEntropyThreshold(float Value) override;
	virtual bool SetSuppressBlank(bool Value) override;
	virtual bool SetSuppressNonSpeechTokens(bool Value) override;
	virtual bool SetBeamSize(int32 Value) override;
	virtual bool SetInitialPrompt(const FString& Value) override;
	virtual bool SetGpuAcceleration(bool Value) override;
	virtual bool SetGpuDeviceId(int32 Value) override;
	//~ End URuntimeSpeechRecognizerProviderBase Interface

	/**
	 * Sets the path to the folder containing the Vosk model to use for the given language
	 *
	 * @param Language The language this model folder corresponds to
	 * @param ModelFolderPath Absolute path to the folder containing the Vosk model (e.g. extracted vosk-model-small-en-us-*)
	 * @return True if the path was set successfully
	 */
	UFUNCTION(BlueprintCallable, Category = "Runtime Speech Recognizer|Vosk")
	bool SetVoskModelPath(ESpeechRecognizerLanguage Language, const FString& ModelFolderPath);

	/**
	 * Sets the model to use for the given language, by the name of a model already known to the
	 * predefined catalog.
	 *
	 * @param Language The language this model corresponds to
	 * @param ModelName Name of a model already present on disk (see URuntimeVoskModelLibrary::GetAvailableVoskModelNames)
	 */
	UFUNCTION(BlueprintCallable, meta = (DisplayName = "Set Vosk Model Path (By Name)"), Category = "Runtime Speech Recognizer|Vosk")
	bool SetVoskModelPathByName(ESpeechRecognizerLanguage Language, UPARAM(Meta = (GetOptions = "RuntimeSpeechRecognizerVosk.RuntimeVoskModelLibrary.GetAvailableVoskModelNames")) const FName ModelName);

protected:
	void Initialize();

private:
	/** The internally owned Vosk recognizer worker thread */
	TSharedPtr<FRuntimeVoskRecognizerThread> Thread;

	/** Model folder paths per language, set via SetVoskModelPath/SetVoskModelPathByName */
	TMap<ESpeechRecognizerLanguage, FString> VoskModelPathsByLanguage;

#if PLATFORM_ANDROID
	TSet<FString> PendingAndroidModelCopies;
#endif

	/** Recognition parameters, mirrored for GetRecognitionParameters()/SetRecognitionParameters() API compatibility */
	FSpeechRecognitionParameters RecognitionParameters;
};