// Georgy Treshchev 2026 All Rights Reserved.

#pragma once

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

struct VoskModel;
struct VoskRecognizer;

/** Static delegate for speech recognition finished recognizing all the queued audio data */
DECLARE_MULTICAST_DELEGATE(FOnVoskRecognitionFinished);

/** Static delegate for recognized words. The recognized text segment is passed as a parameter */
DECLARE_MULTICAST_DELEGATE_OneParam(FOnVoskRecognizedTextSegment, const FString&);

/** Static delegate for speech recognition errors */
DECLARE_MULTICAST_DELEGATE_TwoParams(FOnVoskRecognitionError, const FString& /* ShortErrorMessage */, const FString& /* LongErrorMessage */);

/** Static delegate for speech recognition progress. Vosk has no notion of progress (it's a streaming recognizer, not batch),
 *  so this is only ever broadcast with 100 once queued audio finishes processing, to mirror the whisper.cpp provider's contract */
DECLARE_MULTICAST_DELEGATE_OneParam(FOnVoskRecognitionProgress, int32);

/** Static delegate for the recognizer thread being fully stopped */
DECLARE_MULTICAST_DELEGATE(FOnVoskRecognitionStopped);

/**
 * Vosk API recognizer thread worker. Mirrors FSpeechRecognizerThread's queueing/step-size behavior
 * so that ProcessAudioData(..., bLast) semantics stay identical between providers, but internally
 * drives Vosk's streaming vosk_recognizer_accept_waveform_f API instead of whisper's batch calls
 */
class FRuntimeVoskRecognizerThread : public FRunnable, public TSharedFromThis<FRuntimeVoskRecognizerThread>
{
public:
	FRuntimeVoskRecognizerThread();
	virtual ~FRuntimeVoskRecognizerThread() override;

	/** Sets the path to the model folder to load on the next StartThread() call */
	void SetModelFolderPath(const FString& InModelFolderPath);

	/** Starts the worker thread and loads/references the Vosk model */
	TFuture<bool> StartThread();

	/** Stops the worker thread and dereferences the Vosk model */
	void StopThread();

	/**
	 * Queues audio data for recognition, mirroring FSpeechRecognizerThread::ProcessPCMData's step-size behavior
	 */
	void ProcessPCMData(Audio::FAlignedFloatBuffer PCMData, float SampleRate, uint32 NumOfChannels, bool bLast);

	/** Forces any pending (not yet step-size-filled) audio to be queued for processing */
	void ForceProcessPendingAudioData();

	/** Clears pending/queued audio data */
	void ClearAudioData(bool bClearPendingAudioData, bool bClearAudioQueue);

	//~ Begin FRunnable Interface
	virtual bool Init() override;
	virtual uint32 Run() override;
	virtual void Stop() override;
	virtual void Exit() override;
	//~ End FRunnable Interface

	bool GetIsStopped() const;
	bool GetIsStopping() const;
	bool GetIsFinished() const;

	/** Sets the step size in milliseconds, mirroring FSpeechRecognitionParameters::StepSizeMs. 0 disables step-size batching */
	void SetStepSizeMs(int32 Value);

	/**
	 * Sets the max tokens per segment, mirroring FSpeechRecognitionParameters::MaxTokens.
	 * Vosk itself has no token limit concept, so this is repurposed as a streaming/non-streaming text delivery switch:
	 * - Value > 0 (streaming): recognized text is broadcast via OnRecognizedTextSegment as soon as Vosk produces it
	 * - Value == 0 (non-streaming): recognized text is accumulated internally and only broadcast as a single
	 *   segment once ForceProcessPendingAudioData() or a bLast chunk finalizes the current utterance
	 */
	void SetMaxTokensPerSegment(int32 Value);

	FOnVoskRecognitionFinished OnRecognitionFinished;
	FOnVoskRecognizedTextSegment OnRecognizedTextSegment;
	FOnVoskRecognitionError OnRecognitionError;
	FOnVoskRecognitionProgress OnRecognitionProgress;
	FOnVoskRecognitionStopped OnRecognitionStopped;

private:
	void ReportError(const FString& ShortErrorMessage, const FString& LongErrorMessage);
	void ReleaseMemory();

	/** Parses a Vosk JSON result (partial or final) and either broadcasts it immediately or accumulates it,
	 *  depending on MaxTokensPerSegment */
	void HandleVoskResultJson(const FString& JsonResult, bool bIsPartial);

	/** Broadcasts and clears whatever text has been accumulated so far in non-streaming mode. No-op if empty */
	void FlushAccumulatedText();

private:
	/** A chunk of audio queued for recognition, along with whether it is the last chunk of the current utterance */
	struct FQueuedAudioChunk
	{
		Audio::FAlignedFloatBuffer Buffer;
		bool bLast = false;
	};

	FThreadSafeBool bIsStopped;
	FThreadSafeBool bIsFinished;
	FThreadSafeBool bIsStopping;

	TUniquePtr<FRunnableThread> Thread;

	TQueue<FQueuedAudioChunk> AudioQueue;
	Audio::FAlignedFloatBuffer PendingAudioBuffer;
	FCriticalSection PendingAudioGuard;

	FString ModelFolderPath;
	int32 StepSizeMs = 5000;

	/** 0 = non-streaming (accumulate and flush on finalize), > 0 = streaming (broadcast immediately) */
	int32 MaxTokensPerSegment = 0;

	/** Text accumulated so far in non-streaming mode, awaiting a flush */
	FString AccumulatedText;

	VoskModel* VoskModelPtr;
	VoskRecognizer* VoskRecognizerPtr;

	TUniquePtr<TPromise<bool>> StartThreadPromise;
};