// Georgy Treshchev 2025.

#include "DefaultLipSyncProcessor.h"
#include "ILipSyncProcessor.h"
#include "UObject/WeakObjectPtrTemplates.h"
#include "RuntimeMetaHumanLipSync_Standard.h"

FDefaultLipSyncProcessor::FDefaultLipSyncProcessor()
{
}

FDefaultLipSyncProcessor::~FDefaultLipSyncProcessor()
{
}

void FDefaultLipSyncProcessor::Initialize(UObject* ContextOwner, int32 SampleRate, int32 BufferSize)
{
	TSharedPtr<FLipSyncContextData> ContextData = GetOrCreateContextData(ContextOwner);
	if (!ContextData)
	{
		UE_LOG(LogRuntimeMetaHumanLipSyncStandard, Error, TEXT("Failed to get context data for initialization"));
		return;
	}

#if PLATFORM_WINDOWS || PLATFORM_ANDROID
	ContextData->LipSyncContext = MakeShared<UOVRLipSyncContextWrapper>(
		ovrLipSyncContextProvider_EnhancedWithLaughter,
		SampleRate,
		BufferSize,
		FString(),
		true
	);
#endif

	UE_LOG(LogRuntimeMetaHumanLipSyncStandard, Log, TEXT("DefaultLipSyncProcessor initialized context for owner %s with sample rate %d and buffer size %d"),
	       ContextOwner ? *ContextOwner->GetName() : TEXT("Unknown"), SampleRate, BufferSize);
}

void FDefaultLipSyncProcessor::SetAsyncCallback(UObject* ContextOwner, TFunction<void(const TArray<float>&, float)> Callback)
{
	TSharedPtr<FLipSyncContextData> ContextData = GetOrCreateContextData(ContextOwner);
	if (!ContextData)
	{
		UE_LOG(LogRuntimeMetaHumanLipSyncStandard, Error, TEXT("Failed to get context data for setting callback"));
		return;
	}

	ContextData->CallbackFunction = Callback;

#if PLATFORM_WINDOWS || PLATFORM_ANDROID
	if (ContextData->LipSyncContext)
	{
		ContextData->LipSyncContext->SetAsyncCallback([ContextData](const TArray<float>& NewVisemes, float NewLaughterScore)
		{
			if (ContextData->CallbackFunction)
			{
				ContextData->CallbackFunction(NewVisemes, NewLaughterScore);
			}
		});
	}
#endif
}

void FDefaultLipSyncProcessor::ProcessFrameAsync(UObject* ContextOwner, const float* PCMData, int32 NumSamples)
{
	TSharedPtr<FLipSyncContextData> ContextData = GetOrCreateContextData(ContextOwner);
	if (!ContextData)
	{
		UE_LOG(LogRuntimeMetaHumanLipSyncStandard, Error, TEXT("Failed to get context data for processing frame"));
		return;
	}

#if PLATFORM_WINDOWS || PLATFORM_ANDROID
	if (ContextData->LipSyncContext)
	{
		ContextData->LipSyncContext->ProcessFrameAsync(PCMData, NumSamples);
	}
#endif
}

void FDefaultLipSyncProcessor::CleanupContext(UObject* ContextOwner)
{
	if (!ContextOwner)
	{
		return;
	}

	FScopeLock Lock(&ContextMapMutex);

	TWeakObjectPtr<UObject> WeakOwner(ContextOwner);
	if (ContextMap.Contains(WeakOwner))
	{
		UE_LOG(LogRuntimeMetaHumanLipSyncStandard, Log, TEXT("Cleaning up lip sync context for owner %s"), *ContextOwner->GetName());
		ContextMap.Remove(WeakOwner);
	}
}

TSharedPtr<FDefaultLipSyncProcessor::FLipSyncContextData> FDefaultLipSyncProcessor::GetOrCreateContextData(UObject* ContextOwner)
{
    if (!ContextOwner)
    {
        return nullptr;
    }

    FScopeLock Lock(&ContextMapMutex);

    TWeakObjectPtr<UObject> WeakOwner(ContextOwner);

    for (auto It = ContextMap.CreateIterator(); It; ++It)
    {
        if (!It.Key().IsValid())
        {
            UE_LOG(LogRuntimeMetaHumanLipSyncStandard, Log, TEXT("Cleaning up invalid context from context map"));
            It.RemoveCurrent();
        }
    }

    TSharedPtr<FLipSyncContextData>* ExistingData = ContextMap.Find(WeakOwner);
    if (ExistingData && ExistingData->IsValid())
    {
        return *ExistingData;
    }

    TSharedPtr<FLipSyncContextData> NewContextData = MakeShared<FLipSyncContextData>();
    ContextMap.Add(WeakOwner, NewContextData);

    return NewContextData;
}
