// Georgy Treshchev 2026 All Rights Reserved.

#include "WebRTCAEC3Feature.h"
#include "Codecs/RAW_RuntimeCodec.h"

THIRD_PARTY_INCLUDES_START
#include "aec3_wrapper.h"
THIRD_PARTY_INCLUDES_END

FWebRTCAEC3Feature::FWebRTCAEC3Feature()
	: Instance(nullptr)
	, RenderAudioBuffer(nullptr)
	, CaptureAudioBuffer(nullptr)
	, SampleRate(0)
	, NumChannels(0)
	, FrameSizePerChannel(0)
	, bInitialized(false)
{
}

FWebRTCAEC3Feature::~FWebRTCAEC3Feature()
{
	if (Instance)
	{
		aec3_destroy(Instance);
		Instance = nullptr;
	}
}

bool FWebRTCAEC3Feature::Initialize(int32 InSampleRate, int32 InNumChannels)
{
	FScopeLock Lock(&ProcessingCriticalSection);

	// Validate parameters. WebRTC AEC3 supports 8000, 16000, 32000, 48000 Hz.
	if (InSampleRate != 8000 && InSampleRate != 16000 && InSampleRate != 32000 && InSampleRate != 48000)
	{
		UE_LOG(LogTemp, Error, TEXT("WebRTC AEC3: Unsupported sample rate %d. Supported: 8000, 16000, 32000, 48000"), InSampleRate);
		return false;
	}

	if (InNumChannels < 1)
	{
		UE_LOG(LogTemp, Error, TEXT("WebRTC AEC3: Invalid number of channels: %d"), InNumChannels);
		return false;
	}

	// Reset any existing state
	if (Instance)
	{
		aec3_destroy(Instance);
		Instance = nullptr;
	}
	RenderAudioBuffer = nullptr;
	CaptureAudioBuffer = nullptr;
	AccumulatedRenderData.Reset();
	AccumulatedCaptureData.Reset();

	SampleRate = InSampleRate;
	NumChannels = InNumChannels;
	FrameSizePerChannel = SampleRate / 100; // 10ms frame

	Instance = aec3_create(SampleRate, NumChannels);
	if (!Instance)
	{
		UE_LOG(LogTemp, Error, TEXT("WebRTC AEC3: Failed to create AEC3 instance"));
		bInitialized = false;
		return false;
	}

	// Cache buffer pointers (owned by Instance, valid for its lifetime)
	RenderAudioBuffer = aec3_get_render_buffer(Instance);
	CaptureAudioBuffer = aec3_get_capture_buffer(Instance);

	bInitialized = true;

	UE_LOG(LogTemp, Log, TEXT("WebRTC AEC3: Initialized (SampleRate=%d, NumChannels=%d, FrameSize=%d samples/channel)"),
		SampleRate, NumChannels, FrameSizePerChannel);

	return true;
}

void FWebRTCAEC3Feature::Reset()
{
	FScopeLock Lock(&ProcessingCriticalSection);

	if (Instance)
	{
		aec3_destroy(Instance);
		Instance = nullptr;
	}
	RenderAudioBuffer = nullptr;
	CaptureAudioBuffer = nullptr;
	AccumulatedRenderData.Reset();
	AccumulatedCaptureData.Reset();

	SampleRate = 0;
	NumChannels = 0;
	FrameSizePerChannel = 0;
	bInitialized = false;
}

bool FWebRTCAEC3Feature::IsInitialized() const
{
	return bInitialized;
}

void FWebRTCAEC3Feature::ProcessRenderAudio(const TArray<float>& RenderPCMData, int32 InSampleRate, int32 InNumChannels)
{
	FScopeLock Lock(&ProcessingCriticalSection);

	if (!bInitialized || !Instance)
	{
		return;
	}

	const int32 FrameSizeInterleaved = FrameSizePerChannel * NumChannels;

	// Fast path: no conversion needed
	if (InSampleRate == SampleRate && InNumChannels == NumChannels)
	{
		AccumulatedRenderData.Append(RenderPCMData);

		int32 ConsumedSamples = 0;
		while (AccumulatedRenderData.Num() - ConsumedSamples >= FrameSizeInterleaved)
		{
			ProcessRenderFrame(AccumulatedRenderData.GetData() + ConsumedSamples, FrameSizePerChannel);
			ConsumedSamples += FrameSizeInterleaved;
		}
		if (ConsumedSamples > 0)
		{
#if UE_VERSION_OLDER_THAN(5, 5, 0)
			AccumulatedRenderData.RemoveAt(0, ConsumedSamples, false);
#else
			AccumulatedRenderData.RemoveAt(0, ConsumedSamples, EAllowShrinking::No);
#endif
		}
		return;
	}

	// Slow path: resample and/or remix as needed
	Audio::FAlignedFloatBuffer WorkBuffer(RenderPCMData.GetData(), RenderPCMData.Num());

	if (InSampleRate != SampleRate)
	{
		Audio::FAlignedFloatBuffer ResampledBuffer;
		if (!FRAW_RuntimeCodec::ResampleRAWData(WorkBuffer, InNumChannels, InSampleRate, SampleRate, ResampledBuffer))
		{
			UE_LOG(LogTemp, Error, TEXT("WebRTC AEC3: Failed to resample render audio from %d to %d Hz. Skipping."), InSampleRate, SampleRate);
			return;
		}
		WorkBuffer = MoveTemp(ResampledBuffer);
	}

	if (InNumChannels != NumChannels)
	{
		Audio::FAlignedFloatBuffer RemixedBuffer;
		if (!FRAW_RuntimeCodec::MixChannelsRAWData(WorkBuffer, SampleRate, InNumChannels, NumChannels, RemixedBuffer))
		{
			UE_LOG(LogTemp, Error, TEXT("WebRTC AEC3: Failed to remix render audio from %d to %d channels. Skipping."), InNumChannels, NumChannels);
			return;
		}
		WorkBuffer = MoveTemp(RemixedBuffer);
	}

	AccumulatedRenderData.Append(WorkBuffer.GetData(), WorkBuffer.Num());

	int32 ConsumedSamples = 0;
	while (AccumulatedRenderData.Num() - ConsumedSamples >= FrameSizeInterleaved)
	{
		ProcessRenderFrame(AccumulatedRenderData.GetData() + ConsumedSamples, FrameSizePerChannel);
		ConsumedSamples += FrameSizeInterleaved;
	}
	if (ConsumedSamples > 0)
	{
#if UE_VERSION_OLDER_THAN(5, 5, 0)
		AccumulatedRenderData.RemoveAt(0, ConsumedSamples, false);
#else
		AccumulatedRenderData.RemoveAt(0, ConsumedSamples, EAllowShrinking::No);
#endif
	}
}

bool FWebRTCAEC3Feature::ProcessCaptureAudio(TArray<float>& CapturePCMData, int32 InSampleRate, int32 InNumChannels)
{
	FScopeLock Lock(&ProcessingCriticalSection);

	if (!bInitialized || !Instance)
	{
		return false;
	}

	const int32 FrameSizeInterleaved = FrameSizePerChannel * NumChannels;

	// Fast path: no conversion needed
	if (InSampleRate == SampleRate && InNumChannels == NumChannels)
	{
		AccumulatedCaptureData.Append(CapturePCMData);
	}
	else
	{
		// Slow path: resample and/or remix as needed
		Audio::FAlignedFloatBuffer WorkBuffer(CapturePCMData.GetData(), CapturePCMData.Num());

		if (InSampleRate != SampleRate)
		{
			Audio::FAlignedFloatBuffer ResampledBuffer;
			if (!FRAW_RuntimeCodec::ResampleRAWData(WorkBuffer, InNumChannels, InSampleRate, SampleRate, ResampledBuffer))
			{
				UE_LOG(LogTemp, Error, TEXT("WebRTC AEC3: Failed to resample capture audio from %d to %d Hz. Skipping."), InSampleRate, SampleRate);
				return false;
			}
			WorkBuffer = MoveTemp(ResampledBuffer);
		}

		if (InNumChannels != NumChannels)
		{
			Audio::FAlignedFloatBuffer RemixedBuffer;
			if (!FRAW_RuntimeCodec::MixChannelsRAWData(WorkBuffer, SampleRate, InNumChannels, NumChannels, RemixedBuffer))
			{
				UE_LOG(LogTemp, Error, TEXT("WebRTC AEC3: Failed to remix capture audio from %d to %d channels. Skipping."), InNumChannels, NumChannels);
				return false;
			}
			WorkBuffer = MoveTemp(RemixedBuffer);
		}

		AccumulatedCaptureData.Append(WorkBuffer.GetData(), WorkBuffer.Num());
	}

	TArray<float> ProcessedOutput;
	ProcessedOutput.Reserve(AccumulatedCaptureData.Num());

	int32 ConsumedSamples = 0;
	while (AccumulatedCaptureData.Num() - ConsumedSamples >= FrameSizeInterleaved)
	{
		float* FramePtr = AccumulatedCaptureData.GetData() + ConsumedSamples;
		ProcessCaptureFrame(FramePtr, FrameSizePerChannel);
		ProcessedOutput.Append(FramePtr, FrameSizeInterleaved);
		ConsumedSamples += FrameSizeInterleaved;
	}

	if (ConsumedSamples > 0)
	{
#if UE_VERSION_OLDER_THAN(5, 5, 0)
		AccumulatedCaptureData.RemoveAt(0, ConsumedSamples, false);
#else
		AccumulatedCaptureData.RemoveAt(0, ConsumedSamples, EAllowShrinking::No);
#endif
	}

	if (ProcessedOutput.Num() > 0)
	{
		if (SampleRate != InSampleRate || NumChannels != InNumChannels)
		{
			Audio::FAlignedFloatBuffer WorkBuffer(ProcessedOutput.GetData(), ProcessedOutput.Num());

			if (SampleRate != InSampleRate)
			{
				Audio::FAlignedFloatBuffer ResampledBuffer;
				if (!FRAW_RuntimeCodec::ResampleRAWData(WorkBuffer, NumChannels, SampleRate, InSampleRate, ResampledBuffer))
				{
					UE_LOG(LogTemp, Error, TEXT("WebRTC AEC3: Failed to resample capture output from %d to %d Hz."), SampleRate, InSampleRate);
					return false;
				}
				WorkBuffer = MoveTemp(ResampledBuffer);
			}

			if (NumChannels != InNumChannels)
			{
				Audio::FAlignedFloatBuffer RemixedBuffer;
				if (!FRAW_RuntimeCodec::MixChannelsRAWData(WorkBuffer, InSampleRate, NumChannels, InNumChannels, RemixedBuffer))
				{
					UE_LOG(LogTemp, Error, TEXT("WebRTC AEC3: Failed to remix capture output from %d to %d channels."), NumChannels, InNumChannels);
					return false;
				}
				WorkBuffer = MoveTemp(RemixedBuffer);
			}

			CapturePCMData.Reset();
			CapturePCMData.Append(WorkBuffer.GetData(), WorkBuffer.Num());
		}
		else
		{
			CapturePCMData.Reset();
			CapturePCMData.Append(ProcessedOutput.GetData(), ProcessedOutput.Num());
		}
	}

	return true;
}

void FWebRTCAEC3Feature::SetStreamDelay(int32 DelayMs)
{
	FScopeLock Lock(&ProcessingCriticalSection);
	if (Instance)
	{
		aec3_set_stream_delay(Instance, DelayMs);
	}
}

int32 FWebRTCAEC3Feature::GetSampleRate() const
{
	return SampleRate;
}

int32 FWebRTCAEC3Feature::GetNumChannels() const
{
	return NumChannels;
}

FString FWebRTCAEC3Feature::GetImplementationName() const
{
	return TEXT("WebRTC AEC3");
}

int32 FWebRTCAEC3Feature::GetExpectedFrameSizePerChannel() const
{
	return FrameSizePerChannel;
}

void FWebRTCAEC3Feature::ProcessRenderFrame(const float* FrameData, int32 NumSamplesPerChannel)
{
	// Deinterleave into the render audio buffer
	DeinterleaveToAudioBuffer(FrameData, NumSamplesPerChannel, RenderAudioBuffer);

	// Split into frequency bands for AEC processing
	aec3_buffer_split_into_frequency_bands(RenderAudioBuffer);

	// Analyze the render (far-end) signal
	aec3_analyze_render(Instance);

	// Merge frequency bands back
	aec3_buffer_merge_frequency_bands(RenderAudioBuffer);
}

bool FWebRTCAEC3Feature::ProcessCaptureFrame(float* FrameData, int32 NumSamplesPerChannel)
{
	// Deinterleave into the capture audio buffer
	DeinterleaveToAudioBuffer(FrameData, NumSamplesPerChannel, CaptureAudioBuffer);

	// Analyze the capture signal
	aec3_analyze_capture(Instance);

	// Split into frequency bands
	aec3_buffer_split_into_frequency_bands(CaptureAudioBuffer);

	// Apply high pass filter
	aec3_apply_high_pass_filter(Instance);

	// Process capture to remove echo
	aec3_process_capture(Instance);

	// Merge frequency bands back
	aec3_buffer_merge_frequency_bands(CaptureAudioBuffer);

	// Interleave back into the caller's buffer
	InterleaveFromAudioBuffer(CaptureAudioBuffer, FrameData, NumSamplesPerChannel);

	return true;
}

void FWebRTCAEC3Feature::DeinterleaveToAudioBuffer(const float* InterleavedData, int32 NumSamplesPerChannel, AEC3AudioBuffer* Buffer)
{
    if (!Buffer) return;
    float* const* Channels = aec3_buffer_get_channels(Buffer);

    if (NumChannels == 1)
    {
        if (Channels && Channels[0])
        {
            for (int32 i = 0; i < NumSamplesPerChannel; ++i)
            {
                Channels[0][i] = InterleavedData[i] * 32768.0f;
            }
        }
    }
    else
    {
        if (Channels)
        {
            for (int32 SampleIdx = 0; SampleIdx < NumSamplesPerChannel; ++SampleIdx)
            {
                for (int32 ChannelIdx = 0; ChannelIdx < NumChannels; ++ChannelIdx)
                {
                    if (Channels[ChannelIdx])
                    {
                        Channels[ChannelIdx][SampleIdx] = InterleavedData[SampleIdx * NumChannels + ChannelIdx] * 32768.0f;
                    }
                }
            }
        }
    }
}

void FWebRTCAEC3Feature::InterleaveFromAudioBuffer(const AEC3AudioBuffer* Buffer, float* InterleavedData, int32 NumSamplesPerChannel)
{
    if (!Buffer) return;
    const float* const* Channels = aec3_buffer_get_channels_const(Buffer);

    constexpr float kInvScale = 1.0f / 32768.0f;

    if (NumChannels == 1)
    {
        if (Channels && Channels[0])
        {
            for (int32 i = 0; i < NumSamplesPerChannel; ++i)
            {
                InterleavedData[i] = Channels[0][i] * kInvScale;
            }
        }
    }
    else
    {
        if (Channels)
        {
            for (int32 SampleIdx = 0; SampleIdx < NumSamplesPerChannel; ++SampleIdx)
            {
                for (int32 ChannelIdx = 0; ChannelIdx < NumChannels; ++ChannelIdx)
                {
                    if (Channels[ChannelIdx])
                    {
                        InterleavedData[SampleIdx * NumChannels + ChannelIdx] = Channels[ChannelIdx][SampleIdx] * kInvScale;
                    }
                }
            }
        }
    }
}

