/* Jitter buffering functions
 *
 * (C) 2022 by Andreas Eversberg <jolly@eversberg.eu>
 * All Rights Reserved
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/* How does it work:
 *
 * Storing:
 *
 * Each saved frame is sorted into the list of packages by their sequence
 * number.
 *
 * The first packet will be stored with a delay of minimum jitter window size.
 * 
 * Packets with the same sequence are dropped.
 *
 * Early packts that exceed maximum jitter window size cause jitter
 * window to shift into the future.
 *
 * Late packets cause jitter window to shift into the past (allowing more
 * delay). Minimum jitter window size is added also, to prevent subsequent
 * packets from beeing late too.
 *
 * If no sequence is provided (autosequence), the sequence number is generated
 * by a counter. Also the timestamp is generated by counting the length of each
 * frame.
 *
 * If ssrc changes, the buffer is reset.
 *
 *
 * Playout:
 *
 * The caller of the playout function can request any length of samples from
 * the packet list. The packt's time stamp and the jitter window time stamp
 * indicate what portion of a packet is already provided to the caller.
 * Complete packet, sent to the caller, are removed.
 *
 * Missing packets are interpolated by repeating last 20ms of audio (optional)
 * or by inserting zeroes (sample size > 1 byte) or by inserting 0xff (sample
 * size = 1).
 *
 * Optionally the constant delay will be measured continuously and lowered if
 * greater than minimum window size. (adaptive jitter buffer size)
 *
 * Note that the delay is measured with time stamp of frame, no matter what
 * the length is. Length is an extra delay, but not considered here.
 *
 *
 * Unlocking:
 *
 * If the buffer is created or reset, the buffer is locked, so no packets are
 * stored. When the playout routine is called, the buffer is unlocked. This
 * prevents from filling the buffer before playout is performed, which would
 * cause high delay.
 *
 */

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <math.h>
#include "../libsample/sample.h"
#include "../libdebug/debug.h"
#include "jitter.h"

#define INITIAL_DELAY_INTERVAL	0.5
#define REPEAT_DELAY_INTERVAL	3.0
#define EXTRA_BUFFER		0.020	// 20 ms

/* uncomment to enable heavy debugging */
//#define HEAVY_DEBUG

static int unnamed_count = 1;

/* create jitter buffer */
int jitter_create(jitter_t *jb, const char *name, double samplerate, int sample_size, double target_window_duration, double max_window_duration, uint32_t window_flags)
{
	int rc = 0;
	memset(jb, 0, sizeof(*jb));
	jb->sample_duration = 1.0 / samplerate;
	jb->sample_size = sample_size;
	jb->target_window_size = (int)(samplerate * target_window_duration);
	jb->max_window_size = (int)(samplerate * max_window_duration);
	jb->window_flags = window_flags;

	jb->extra_size = (int)(EXTRA_BUFFER * samplerate);
	jb->extra_samples = calloc(sample_size, jb->extra_size);
	if (!jb->extra_samples) {
		PDEBUG(DJITTER, DEBUG_ERROR, "No memory for frame.\n");
		rc = -ENOMEM;
		goto error;
	}


	/* optionally give a string to be show with the debug */
	if (name && *name)
		snprintf(jb->name, sizeof(jb->name) - 1, "(%s) ", name);
	else
		snprintf(jb->name, sizeof(jb->name) - 1, "(unnamed %d) ", unnamed_count++);

	jitter_reset(jb);

	PDEBUG(DJITTER, DEBUG_INFO, "%sCreated jitter buffer. (samplerate=%.0f, target_window=%.0fms, max_window=%.0fms, flag:latency=%s flag:repeat=%s)\n", jb->name, samplerate, target_window_duration * 1000.0, max_window_duration * 1000.0, (window_flags & JITTER_FLAG_LATENCY) ? "true" : "false", (window_flags & JITTER_FLAG_REPEAT) ? "true" : "false");

error:
	if (rc)
		jitter_destroy(jb);
	return rc;
}

/* reset jitter buffer */
void jitter_reset(jitter_t *jb)
{
	jitter_frame_t *jf, *temp;

	PDEBUG(DJITTER, DEBUG_INFO, "%sReset jitter buffer.\n", jb->name);

	/* jitter buffer locked */
	jb->unlocked = 0;

	/* window becomes invalid */
	jb->window_valid = 0;

	/* remove all pending frames */
	jf = jb->frame_list;
	while(jf) {
		temp = jf;
		jf = jf->next;
		free(temp);
	}
	jb->frame_list = NULL;

	/* clear extrapolation buffer */
	if (jb->extra_samples) {
		if (jb->sample_size == 1)
			memset(jb->extra_samples, 0xff, jb->sample_size * jb->extra_size);
		else
			memset(jb->extra_samples, 0, jb->sample_size * jb->extra_size);
	}
	jb->extra_index = 0;

	/* delay measurement and reduction */
	jb->delay_counter = 0.0;
	jb->delay_interval = INITIAL_DELAY_INTERVAL;
	jb->min_delay_value = -1;
}

void jitter_destroy(jitter_t *jb)
{
	jitter_reset(jb);

	PDEBUG(DJITTER, DEBUG_INFO, "%sDestroying jitter buffer.\n", jb->name);

	if (jb->extra_samples) {
		free(jb->extra_samples);
		jb->extra_samples = NULL;
	}
}

/* store audio in jitterbuffer
 *
 * stop if buffer is completely filled
 */
void jitter_save(jitter_t *jb, void *samples, int length, int has_sequence, uint16_t sequence, uint32_t timestamp, uint32_t ssrc)
{
	jitter_frame_t *jf, **jfp;
	int16_t offset_sequence;
	int32_t offset_timestamp;

	/* ignore frames until the buffer is unlocked by jitter_load() */
	if (!jb->unlocked)
		return;

	/* omit frames with no data */
	if (length < 1)
		return;

	/* generate sequence and timestamp automatically, if enabled */
	if (!has_sequence) {
#ifdef DEBUG_JITTER
		PDEBUG(DJITTER, DEBUG_DEBUG, "%sSave frame of %d samples (no seqence).\n", jb->name, length);
#endif
		sequence = jb->next_sequence;
		jb->next_sequence++;
		timestamp = jb->next_timestamp;
		jb->next_timestamp += length;
		ssrc = jb->window_ssrc;
	} else {
#ifdef HEAVY_DEBUG
		PDEBUG(DJITTER, DEBUG_DEBUG, "%sSave frame of %d samples (seqence=%u timestamp=%u ssrc=0x%02x).\n", jb->name, length, sequence, timestamp, ssrc);
#endif
		jb->next_sequence = sequence + 1;
		jb->next_timestamp = timestamp + length;
	}

	/* first packet (with this ssrc) sets window size to target_window_size */
	if (!jb->window_valid || jb->window_ssrc != ssrc) {
		if (!jb->window_valid)
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s Initial frame after init or reset.\n", jb->name);
		else
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s SSRC changed.\n", jb->name);
		// NOTE: Reset must be called before finding the frame location below, because there will be no frame in list anymore!
		jitter_reset(jb);
		jb->unlocked = 1;
		/* when using dynamic jitter buffer, we use half of the target delay */
		if ((jb->window_flags & JITTER_FLAG_LATENCY)) {
			jb->window_timestamp = timestamp - (uint32_t)jb->target_window_size / 2;
		} else {
			jb->window_timestamp = timestamp - (uint32_t)jb->target_window_size;
	       	}
		jb->window_valid = 1;
		jb->window_ssrc = ssrc;
	}

	/* find location where to put frame into the list, depending on sequence number */
	jfp = &jb->frame_list;
	while(*jfp) {
		offset_sequence = (int16_t)(sequence - (*jfp)->sequence);
		/* found double entry */
		if (offset_sequence == 0) {
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s Dropping double packet (sequence = %d)\n", jb->name, sequence);
			return;
		}
		/* offset is negative, so we found the position to insert frame */
		if (offset_sequence < 0)
			break;
		jfp = &((*jfp)->next);
	}

	offset_timestamp = timestamp - jb->window_timestamp;
#ifdef HEAVY_DEBUG
	PDEBUG(DJITTER, DEBUG_DEBUG, "%sFrame has offset of %.0fms in jitter buffer.\n", jb->name, (double)offset_timestamp * jb->sample_duration * 1000.0);
#endif

	/* measure delay */
	if (jb->min_delay_value < 0 || offset_timestamp < jb->min_delay_value)
		jb->min_delay_value = offset_timestamp;

	/* if frame is too early (delay ceases), shift window to the future */
	if (offset_timestamp > jb->max_window_size) {
		if ((jb->window_flags & JITTER_FLAG_LATENCY)) {
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s Frame too early: Shift jitter buffer to the future, to make the frame fit to the end. (offset_timestamp(%d) > max_window_size(%d))\n", jb->name, offset_timestamp, jb->max_window_size);
			/* shift window so it fits to the end of window */
			jb->window_timestamp = timestamp - jb->max_window_size;
		} else {
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s Frame too early: Shift jitter buffer to the future, to make the frame fit to the target delay. (offset_timestamp(%d) > max_window_size(%d))\n", jb->name, offset_timestamp, jb->max_window_size);
			/* shift window so frame fits to the start of window + target delay */
			jb->window_timestamp = timestamp - (uint32_t)(jb->target_window_size);
		}
	}

	/* is frame is too late, shift window to the past. */
	if (offset_timestamp < 0) {
		if ((jb->window_flags & JITTER_FLAG_LATENCY)) {
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s Frame too late: Shift jitter buffer to the past, and add target window size. (offset_timestamp(%d) < 0)\n", jb->name, offset_timestamp);
			/* shift window so frame fits to the start of window + half of target delay */
			jb->window_timestamp = timestamp - (uint32_t)(jb->target_window_size) / 2;
		} else {
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s Frame too late: Shift jitter buffer to the past, and add half target window size. (offset_timestamp(%d) < 0)\n", jb->name, offset_timestamp);
			/* shift window so frame fits to the start of window + target delay */
			jb->window_timestamp = timestamp - (uint32_t)(jb->target_window_size);
		}
	}

	/* insert or append frame */
#ifdef HEAVY_DEBUG
	PDEBUG(DJITTER, DEBUG_DEBUG, "%s Store frame\n", jb->name);
#endif
	jf = malloc(sizeof(*jf) + length * jb->sample_size);
	if (!jf) {
		PDEBUG(DJITTER, DEBUG_ERROR, "No memory for frame.\n");
		return;
	}
	memset(jf, 0, sizeof(*jf)); // note: clear header only
	jf->sequence = sequence;
	jf->timestamp = timestamp;
	memcpy(jf->samples, samples, length * jb->sample_size);
	jf->length = length;
	jf->next = *jfp;
	*jfp = jf;
}

/* get audio from jitterbuffer
 */
void jitter_load(jitter_t *jb, void *samples, int length)
{
	jitter_frame_t *jf;
	int32_t count, count2, index;

#ifdef HEAVY_DEBUG
	PDEBUG(DJITTER, DEBUG_DEBUG, "%sLoad chunk of %d samples.\n", jb->name, length);
#endif

	/* now unlock jitter buffer */
	jb->unlocked = 1;

	/* reduce delay */
	jb->delay_counter += jb->sample_duration * (double)length;
	if (jb->delay_counter >= jb->delay_interval) {
		if (jb->min_delay_value >= 0)
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s Statistics: target_window_delay=%.0fms max_window_delay=%.0fms  current min_delay=%.0fms\n", jb->name, (double)jb->target_window_size * jb->sample_duration * 1000.0, (double)jb->max_window_size * jb->sample_duration * 1000.0, (double)jb->min_delay_value * jb->sample_duration * 1000.0);
		/* delay reduction, if maximum delay is greater than target jitter window size */
		if ((jb->window_flags & JITTER_FLAG_LATENCY) && jb->min_delay_value > jb->target_window_size) {
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s Reducing current minimum delay of %.0fms, because maximum delay is greater than target window size of %.0fms.\n", jb->name, (double)jb->min_delay_value * jb->sample_duration * 1000.0, (double)jb->target_window_size * jb->sample_duration * 1000.0);
			/* only reduce delay to half of the target window size */
			jb->window_timestamp += jb->min_delay_value - jb->target_window_size / 2;

		}
		jb->delay_counter -= jb->delay_interval;
		jb->delay_interval = REPEAT_DELAY_INTERVAL;
		jb->min_delay_value = -1;
	}

	/* process all frames until output buffer is loaded */
	while (length) {
		/* always get frame with the lowest sequence number (1st frame) */
		jf = jb->frame_list;

		if (jf) {
			count = jf->timestamp - jb->window_timestamp;
			if (count > length)
				count = length;
		} else
			count = length;
		/* if there is no frame or we have not reached frame's time stamp, extrapolate */
		if (count > 0) {
#ifdef HEAVY_DEBUG
			if (jf)
				PDEBUG(DJITTER, DEBUG_DEBUG, "%s There is a frame ahead in buffer after %d samples. Interpolating gap.\n", jb->name, jf->timestamp - jb->window_timestamp);
			else
				PDEBUG(DJITTER, DEBUG_DEBUG, "%s There is no frame ahead in buffer. Interpolating gap.\n", jb->name);
#endif
			/* extrapolate by playing the extrapolation buffer */
			while (count) {
				count2 = count;
				if (count2 > jb->extra_size - jb->extra_index)
					count2 = jb->extra_size - jb->extra_index;
				memcpy(samples, (uint8_t *)jb->extra_samples + jb->extra_index * jb->sample_size, count2 * jb->sample_size);
				jb->extra_index += count2;
				if (jb->extra_index == jb->extra_size)
					jb->extra_index = 0;
				samples = (uint8_t *)samples + count2 * jb->sample_size;
				length -= count2;
				jb->window_timestamp += count2;
				count -= count2;
			}
			if (length == 0)
				return;
		}

		/* copy samples from frame (what is not in the past) */
		index = jb->window_timestamp - jf->timestamp;
		while (index < jf->length) {
			/* use the lowest value of 'playout length' or 'remaining packet length' */
			count = length;
			if (jf->length - index < count)
				count = jf->length - index;
			/* if extrapolation is used, limit count to what we can store into buffer */
			if (jb->extra_samples && jb->extra_size - jb->extra_index < count)
				count = jb->extra_size - jb->extra_index;
			/* copy samples from packet to play out, increment sample pointer and decrement length */
#ifdef HEAVY_DEBUG
			PDEBUG(DJITTER, DEBUG_DEBUG, "%s Copy data (offset=%u count=%u) from frame (sequence=%u timestamp=%u length=%u).\n", jb->name, index, count, jf->sequence, jf->timestamp, jf->length);
#endif
			memcpy(samples, (uint8_t *)jf->samples + index * jb->sample_size, count * jb->sample_size);
			samples = (uint8_t *)samples + count * jb->sample_size;
			length -= count;
			/* copy frame data to extrapolation buffer also, increment index */
			if ((jb->window_flags & JITTER_FLAG_REPEAT)) {
				memcpy((uint8_t *)jb->extra_samples + jb->extra_index * jb->sample_size, (uint8_t *)jf->samples + index * jb->sample_size, count * jb->sample_size);
				jb->extra_index += count;
				if (jb->extra_index == jb->extra_size)
					jb->extra_index = 0;
			}
			/* increment time stamp */
			jb->window_timestamp += count;
			index += count;
			/* if there was enough to play out, we are done */
			if (length == 0)
				return;
		}

		/* free frame, because all samples are now in the past */
		jb->frame_list = jf->next;
		free(jf);

		/* now go for next loop, in case there is still date to play out */
	}
}