#include "micpath.h"
#include <cstdio>

aec_stream::aec_stream(pw_loop* lp,ringbuffer* hp,ringbuffer* mic,
	logger_record& ps, aec_params p):
	pwstream(48000, 1, 1, lp, "aec-output", 
		pw_properties_new(PW_KEY_MEDIA_TYPE, "Audio",
		                  PW_KEY_MEDIA_CATEGORY, "Source",
		                  PW_KEY_MEDIA_ROLE, "DSP", NULL)),
	hp_rb(hp), mic_rb(mic), pstrm(ps), params(p) {
	ctx = dsp_EchoCancel_ctx_new_ctx();
}

void aec_stream::on_process() {
	pw_buffer* pwb;
	if((pwb = pw_stream_dequeue_buffer(this->strm)) == nullptr) {
		return;
	}

	auto* b = pwb->buffer;
	float* dst = (float*)b->datas[0].data;
	if(dst == nullptr) return;

	auto n_frames = b->datas[0].maxsize / sizeof(float);
	if(pwb->requested) {
		n_frames = SPA_MIN(pwb->requested, n_frames);
	}

	size_t effective = 0;

	switch(status) {
		case INIT:
			status = SKIP_EARLY;
			remaining_samples = 48000*10;
			[[fallthrough]];
		case SKIP_EARLY:
			effective = mic_rb->dequeue(dst, n_frames, true, false);
			if(remaining_samples >= 48000 * 2 && remaining_samples - effective < 48000*2) {
				std::puts("Please keep silent for the next few seconds");
			}
			if((remaining_samples-=effective) <= 0) {
				status = SCAN_SILENT;
				prm_mutex.lock();
				remaining_samples = 48000;
			}
			break;

		case SCAN_SILENT:
			effective = mic_rb->dequeue(dst, n_frames, true, false);
			for(size_t i = 0; i < effective; ++i) {
				params.noise_power += dst[i] * dst[i];
			}
			if((remaining_samples-=effective) <= 0) {
				status = SCAN_LOUD;
				params.noise_power /= 48000 - remaining_samples;
				remaining_samples = 48000;
				pstrm.request_noise(true, .1);
			}
			break;

		case SCAN_LOUD: {
			effective = mic_rb->dequeue(dst, n_frames, true, false);

			for(size_t i = 1; i <= effective; ++i) {
				for(size_t j = 0; j < 8192; ++j) {
					int t = hp_rb->r_head_pos - 2*i - 2*j;

					while(t < 0) {t += ringbuffer::bsize;}
					if(t >= ringbuffer::bsize) {t -= ringbuffer::bsize;}

					/* -i with (r_head_pos - i) - j*/
					cov[j] += dst[effective - i] * hp_rb->buffer[t + 1];
				}
				params.echo_return_gain += dst[effective-i]*dst[effective-i];	
			}

			if((remaining_samples-=effective) <= 0) {
				status = ACTIVE;
				params.echo_return_gain /= 48000 - remaining_samples;
				params.echo_return_gain -= params.noise_power;
				params.echo_return_gain *= 1/(.1*.1);
				// Renormalize

				float max_cov2 = 0.0f;
				for(int i = 0; i < 8192; ++i) {
					if(cov[i]*cov[i] >= max_cov2) {
						max_cov2 = cov[i]*cov[i];
						params.sample_delay = i;
					}
				}

				params.valid = true;
				remaining_samples = 48000 * 5;
				prm_mutex.unlock();
			} break; }

		case ACTIVE: {
			effective = mic_rb->dequeue(dst, n_frames, true, false);
			// -i with r_head_pos - i - (sample_delay + 8)
			int t = hp_rb->r_head_pos - 2*params.sample_delay + 16 - 2*effective;
			while(t < 0) {t += ringbuffer::bsize;}
			if(t >= ringbuffer::bsize) {t -= ringbuffer::bsize;}

			for(size_t i = effective; i >= 1; --i) {
				double dsp_out;
				dsp_EchoCancel_step(hp_rb->buffer[t+1],
				                    (double)dst[effective - i],
				                    (remaining_samples >= 0) ? 1.0/256 : 1e-4/256,
				                    (double)params.noise_power,
				                    remaining_samples >= 0,
														&dsp_out,
														h, &cnt, ctx);
				dst[effective - i] = dsp_out;

				t += 2;
				if(t == ringbuffer::bsize) {t = 0;}
			}

			if((remaining_samples -= effective) <= 0) {
				pstrm.request_noise(false, 0.f);
			}

			break;
		}

		case INACTIVE:
			effective = mic_rb->dequeue(dst, n_frames, true, false);
			break;
	}

	b->datas[0].chunk->offset = 0;
	b->datas[0].chunk->stride = sizeof(float);
	b->datas[0].chunk->size = effective*sizeof(float);
	pw_stream_queue_buffer(this->strm, pwb);
}

aec_params aec_stream::get_aec_params() {
	prm_mutex.lock();
	auto copy = params;
	prm_mutex.unlock();
	return copy;
}