---

#include "gst_imx_ai_aecnr_core.h"
#include "mediastreamer2/msfilter.h"
#include "mediastreamer2/flowcontrol.h"

typedef struct ImxAiAecState {
	struct imx_ai_aecnr_core *aec_core;
	MSBufferizer delayed_ref;
	MSFlowControlledBufferizer ref;
	MSBufferizer echo;
	int samplerate;
	int framesize;
	int delay_ms;
	int nominal_ref_samples;
	bool_t echostarted;
	bool_t bypass_mode;
	bool_t using_zeroes;
} ImxAiAecState;

static void imx_ai_aec_init(MSFilter *f) {
	ImxAiAecState *s = ms_new0(ImxAiAecState, 1);
	s->samplerate = AI_AECNR_SAMPLE_RATE; // 16000
	s->framesize = AI_AECNR_IO_SAMPLES;   // 256
	
	ms_bufferizer_init(&s->delayed_ref);
	ms_bufferizer_init(&s->echo);
	ms_flow_controlled_bufferizer_init(&s->ref, f, s->samplerate, 1);
	
	s->delay_ms = 0;
	s->echostarted = FALSE;
	s->bypass_mode = FALSE;
	s->using_zeroes = FALSE;
	
	f->data = s;
}

static void imx_ai_aec_preprocess(MSFilter *f) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	imx_ai_aecnr_config config;
	int delay_samples;
	mblk_t *m;

	config.use_small_model = 0; 
	s->aec_core = imx_ai_aecnr_core_open(&config);
	
	if (s->aec_core == NULL) {
		ms_error("imx_ai_aecnr_core_open() failed, bypass mode");
		s->bypass_mode = TRUE;
		return;
	}

	s->echostarted = FALSE;
	delay_samples = s->delay_ms * s->samplerate / 1000;
	s->nominal_ref_samples = delay_samples;

	ms_message("IMX AI AECNR initialized: %s, framesize=%i, delay=%i ms", 
				config.version, s->framesize, s->delay_ms);

	// Настройка flow control для синхронизации потоков
	ms_flow_controlled_bufferizer_set_samplerate(&s->ref, s->samplerate);
	ms_flow_controlled_bufferizer_set_max_size_ms(&s->ref, s->delay_ms);

	/* Заполнение начальной задержки нулями */
	m = allocb(delay_samples * 2, 0);
	memset(m->b_wptr, 0, delay_samples * 2);
	m->b_wptr += delay_samples * 2;
	ms_bufferizer_put(&s->delayed_ref, m);
}

static void imx_ai_aec_process(MSFilter *f) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	int nbytes = s->framesize * sizeof(int16_t);
	mblk_t *refm;
	int16_t *ref_tmp = (int16_t *)alloca(nbytes);
	int16_t *echo_tmp = (int16_t *)alloca(nbytes);

	if (s->bypass_mode) {
		while ((refm = ms_queue_get(f->inputs[0])) != NULL) ms_queue_put(f->outputs[0], refm);
		while ((refm = ms_queue_get(f->inputs[1])) != NULL) ms_queue_put(f->outputs[1], refm);
		return;
	}

	// 1. Получаем референсный сигнал (от дальнего конца)
	if (f->inputs[0] != NULL) {
		if (s->echostarted) {
			while ((refm = ms_queue_get(f->inputs[0])) != NULL) {
				ms_bufferizer_put(&s->delayed_ref, dupmsg(refm));
				ms_flow_controlled_bufferizer_put(&s->ref, refm);
			}
		} else {
			ms_queue_flush(f->inputs[0]);
		}
	}

	// 2. Получаем сигнал с микрофона
	ms_bufferizer_put_from_queue(&s->echo, f->inputs[1]);

	// 3. Основной цикл обработки
	while (ms_bufferizer_get_avail(&s->echo) >= (size_t)nbytes) {
		mblk_t *out_clean = allocb(nbytes, 0);
		int avail;

		if (!s->echostarted) s->echostarted = TRUE;

		// Проверка наличия данных в буфере задержки
		avail = (int)ms_bufferizer_get_avail(&s->delayed_ref);
		if (avail < ((s->nominal_ref_samples * 2) + nbytes)) {
			// Инъекция тишины, если референс отстал
			refm = allocb(nbytes, 0);
			memset(refm->b_wptr, 0, nbytes);
			refm->b_wptr += nbytes;
			ms_bufferizer_put(&s->delayed_ref, refm);
			// В отличие от Speex, тут мы не пушим тишину в output[0], 
			// чтобы не плодить лишнюю задержку в звуковой карте
			if (!s->using_zeroes) {
				ms_warning("IMX AEC: Not enough ref samples, using zeroes");
				s->using_zeroes = TRUE;
			}
		} else {
			s->using_zeroes = FALSE;
			refm = allocb(nbytes, 0);
			ms_flow_controlled_bufferizer_read(&s->ref, refm->b_wptr, nbytes);
			refm->b_wptr += nbytes;
			ms_queue_put(f->outputs[0], refm);
		}

		// Читаем выровненные данные
		ms_bufferizer_read(&s->echo, (uint8_t *)echo_tmp, nbytes);
		ms_bufferizer_read(&s->delayed_ref, (uint8_t *)ref_tmp, nbytes);

		// AI Обработка (AEC + NR)
		if (imx_ai_aecnr_core_process(s->aec_core, ref_tmp, echo_tmp, (int16_t *)out_clean->b_wptr) != 0) {
			ms_error("imx_ai_aecnr_core_process failed");
			memcpy(out_clean->b_wptr, echo_tmp, nbytes); // Fallback
		}

		out_clean->b_wptr += nbytes;
		ms_queue_put(f->outputs[1], out_clean);
	}
}

static void imx_ai_aec_postprocess(MSFilter *f) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	if (s->aec_core) {
		imx_ai_aecnr_core_close(s->aec_core);
		s->aec_core = NULL;
	}
	ms_bufferizer_flush(&s->delayed_ref);
	ms_bufferizer_flush(&s->echo);
}

static void imx_ai_aec_uninit(MSFilter *f) {
	ImxAiAecState *s = (ImxAiAecState *)f->data;
	ms_bufferizer_uninit(&s->delayed_ref);
	ms_bufferizer_uninit(&s->echo);
	ms_free(s);
}