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

#pragma once

#include "CoreMinimal.h"
#include "Components/ActorComponent.h"

THIRD_PARTY_INCLUDES_START
#include "pocketsphinx.h" 
THIRD_PARTY_INCLUDES_END
#include "Containers/Queue.h"
#include "InworldRuntimeTypes.h"
#include "HAL/Runnable.h"
#include "HAL/RunnableThread.h"
#include "InworldLipsyncComponent.generated.h"


struct FInworldData_TTSOutput;

// Forward declarations
class UInworldLipsyncComponent;

class FAudioProcessingRunnable : public FRunnable
{
public:
	// Constructor
	FAudioProcessingRunnable(UInworldLipsyncComponent* InOwnerComponent);

	// Destructor
	virtual ~FAudioProcessingRunnable();

	// FRunnable interface
	virtual uint32 Run() override;
	virtual void Stop() override;
	virtual void Exit() override;

	/** Signals the thread to stop running. */
	void EnsureCompletion();
	void ResetBuffer();

	void InitializePocketSphinx();
	void ShutdownPocketSphinx();

	//On voice complete, make sure we output the empty viseme.
	void VoiceCompleted();
private:
	// PocketSphinx required sample rate.
	const int32 PocketSphinxSampleRate = 16000;
	ps_config_t* SphinxConfig = nullptr;
	ps_decoder_t* SphinxDecoder = nullptr;

	// Reference to the owner component to access its data
	TWeakObjectPtr<UInworldLipsyncComponent> OwnerComponent;

	// Thread control
	std::atomic<bool> bStopping;
	std::atomic<bool> bVoiceComplete;
	FRunnableThread* Thread;

	bool bIsPocketSphinxCreated = false;

	// Local buffer for accumulating audio before processing
	TArray<int16> ProcessingBuffer;
};

// Viseme data coming from PocketSphinx.
UENUM(BlueprintType)
enum class EViseme : uint8
{
	VS_PP UMETA(DisplayName = "PP"),	// p, b, m
	VS_FF UMETA(DisplayName = "FF"),	// f, v
	VS_TH UMETA(DisplayName = "TH"),	// th
	VS_DD UMETA(DisplayName = "DD"),	// t, d
	VS_kk UMETA(DisplayName = "kk"),	// k, g, n
	VS_CH UMETA(DisplayName = "CH"),	// ch, j, sh
	VS_SS UMETA(DisplayName = "SS"),	// s, z
	VS_nn UMETA(DisplayName = "nn"),	// n, l
	VS_RR UMETA(DisplayName = "RR"),	// r
	VS_aa UMETA(DisplayName = "aa"),	// a, ah
	VS_E UMETA(DisplayName = "E"),		// e, eh
	VS_I UMETA(DisplayName = "I"),		// ih
	VS_O UMETA(DisplayName = "O"),		// o, ow
	VS_U UMETA(DisplayName = "U"),		// u, oo
	VS_Silence UMETA(DisplayName = "Silence")
};

DECLARE_DYNAMIC_MULTICAST_DELEGATE_OneParam(FOnInworldVisemeUpdate, const FInworldVisemeBlends&, VisemeBlend);

UCLASS(ClassGroup = (Custom), meta = (BlueprintSpawnableComponent))
class INWORLDLIPSYNC_API UInworldLipsyncComponent : public UActorComponent
{
	GENERATED_BODY()

public:
	UPROPERTY(BlueprintAssignable)
	FOnInworldVisemeUpdate OnVisemeUpdate;

	UInworldLipsyncComponent();

	virtual void BeginPlay() override;
	virtual void EndPlay(const EEndPlayReason::Type EndPlayReason) override;

	void ResetBuffer();

	EViseme MapPhonemeToViseme(const FString& Phoneme);

	// Downsamples the audio to TargetSampleRate. (PocketSphinx requires 16Khz)
	bool DownsampleWaveform(TArray<int16>& Waveform, int32 CurrentSampleRate, int32 TargetSampleRate);
	void UpdateVisemeFromWorker(EViseme NewViseme);
	TQueue<TArray<int16>>& GetAudioQueue()
	{
		return AudioQueue;
	}

	int32 GetSourceSampleRate() const
	{
		return SourceSampleRate.load(std::memory_order_relaxed);
	}

protected:
	void InitializePocketSphinx();
	void ShutdownPocketSphinx();

	// Used by delegate on AudioThread from InworldVoiceAudiotComponent
	void ProcessAudioFrame(TArray<int16>&& AudioData, const int32 InSourceSampleRate);

	// Resamples the audio to TargetSampleRate. (PocketSphinx requires 16Khz) (COSTLY!)
	bool ResampleWaveform(TArray<int16>& Waveform, int32 CurrentSampleRate, int32 TargetSampleRate);
	bool Resample24kTo16k(TArray<int16>& Waveform);

	// Used by delegate on GameThread from InworldVoiceAudiotComponent
	void OnVoiceComplete(class UInworldVoiceAudioComponent* AudioComponent, const FInworldData_TTSOutput& Voice);

	TMap<FString, EViseme> PhonemeToVisemeMap;
	EViseme CurrentViseme = EViseme::VS_Silence;

	std::atomic<int32> SourceSampleRate;
	// Thread safe way to store audio data, without requiring a lock.
	TQueue<TArray<int16>> AudioQueue;

	FDelegateHandle AudioDelegateHandle;

	/** A unique pointer to our runnable worker. This manages its lifetime automatically. */
	TUniquePtr<FAudioProcessingRunnable> AudioProcessor;
};
