/*
 * dgrambuf.c - C datagrams buffer.
 * 
 * Copyright 2016 by Ludovic Pouzenc <ludovic@pouzenc.fr>
 */
#define _GNU_SOURCE /* See feature_test_macros(7) */

#include "dgrambuf.h"

#include <sys/socket.h> /* recvmmsg() _GNU_SOURCE */
#include <stdlib.h> /* calloc(), free(), qsort() */
#include <stdio.h> /* perror() */
#include <string.h> /* memset() */
#include <sys/uio.h> /* writev() */

struct uint_pair {
	unsigned int index;
	unsigned int value;
};

struct dgrambuf_t {
	size_t dgram_count;
	size_t dgram_max_size;
	size_t dgram_header_size;

	struct iovec *recv_iovecs;
	struct iovec *write_iovecs;
	struct mmsghdr *msgs;

	unsigned int win_base;
	unsigned int *dgram_seq_numbers; /* Stores the decoded datagram sequence number for each dgram slot of buf */
	unsigned int *dgram_len;
	struct uint_pair *dgram_ordered_seq_numbers;

	void *buf;

	unsigned int (*validate_func)(unsigned int, void *);
	//TODO pthread_mutex_lock
};

int _compare_uint_pair(const void *pa, const void *pb);

void dgrambuf_set_validate_func(dgrambuf_t dbuf, unsigned int (*func)(unsigned int, void *) ) {
	dbuf->validate_func = func;
}

int dgrambuf_recvmmsg(dgrambuf_t dbuf, int sockfd) {
  void *dgram_base;
	size_t vlen, i, dgram_index;
	int recv_msg_count, res;
	unsigned int seq, dgram_len;

	if ( !dbuf->validate_func ) {
		return -1;
	}

	/* Initialize recvmmsg() syscall arguments */
	for (i=0, vlen=0; i < dbuf->dgram_count; i++) {
		if ( dbuf->dgram_seq_numbers[i] == 0 ) {
			dbuf->recv_iovecs[vlen].iov_base = dbuf->buf + i*dbuf->dgram_max_size;
			dbuf->recv_iovecs[vlen].iov_len = dbuf->dgram_max_size;
			memset(dbuf->msgs + vlen, 0, sizeof(struct mmsghdr));
			dbuf->msgs[vlen].msg_hdr.msg_iov = dbuf->recv_iovecs + vlen;
			dbuf->msgs[vlen].msg_hdr.msg_iovlen = 1;
			vlen++;
		}
	}

	/* Buffer is full, can't receive */
	if ( vlen==0 ) {
		return -2;
	}

	/* Do the syscall */
	recv_msg_count = recvmmsg(sockfd, dbuf->msgs, vlen, MSG_WAITFORONE, NULL);
	if (recv_msg_count < 0) {
		perror("recvmmsg()");
		return recv_msg_count;
	}

	/* Check all received messages */
	res = 1;
	for (i=0; i<recv_msg_count; i++) {
		dgram_base = dbuf->recv_iovecs[i].iov_base;
		dgram_index = (dgram_base - dbuf->buf) / dbuf->dgram_max_size;
		dgram_len = dbuf->msgs[i].msg_len;
		seq = dbuf->validate_func(dgram_len, dgram_base);
		// TODO better feedback
		if ( seq == 0 ) {
			printf("#%zi invalid (%u)\n", i, seq);
			dbuf->dgram_seq_numbers[dgram_index] = 0;
		} else if ( seq == -1 ) {
			printf("#%zi end\n", i);
			dbuf->dgram_seq_numbers[dgram_index] = 0;
			res = 0;
		} else if ( seq < dbuf->win_base ) {
			printf("#%zi past (%u)\n", i, seq);
			dbuf->dgram_seq_numbers[dgram_index] = 0;
		} else if ( seq >= dbuf->win_base + dbuf->dgram_count ) {
			printf("#%zi future (%u)\n", i, seq);
			dbuf->dgram_seq_numbers[dgram_index] = 0;
		} else {
			printf("#%zi valid (%u)\n", i, seq);
			dbuf->dgram_seq_numbers[dgram_index] = seq;
			dbuf->dgram_len[dgram_index] = dgram_len;
		}
	}

	return res;
}

ssize_t dgrambuf_write(dgrambuf_t dbuf, int fd) {
	size_t dgram_index, i, vlen;
	unsigned int curr_seq, prev_seq, dgram_len;
	ssize_t nwrite, total;

	/* Initialize dgram_ordered_seq_numbers from dgram_seq_numbers */
	for (i=0; i < dbuf->dgram_count; i++) {
		dbuf->dgram_ordered_seq_numbers[i].index = i;
		dbuf->dgram_ordered_seq_numbers[i].value = dbuf->dgram_seq_numbers[i];
	}
	/* Inplace sorting of dgram_ordered_seq_numbers */
	qsort(dbuf->dgram_ordered_seq_numbers, dbuf->dgram_count, sizeof(struct uint_pair), _compare_uint_pair);
	
	/* Initialize iovecs for writev, take dgram payloads following the sequence numbers */
	for (prev_seq=0, vlen=0, total=0, i=0; i< dbuf->dgram_count; i++) {
		curr_seq = dbuf->dgram_ordered_seq_numbers[i].value;

		/* Skip empty dgram slot */
		if ( curr_seq == 0 )
			continue;

		/* Skip dgram comming from the past */
		if ( curr_seq < dbuf->win_base ) {
			fprintf(stderr, "Oops : found dgram from past in buffer (%u)\n", curr_seq);
			continue;
		}

		/* Break if first dgram to write is not in buffer at all */
		if ( ( vlen==0 ) && (curr_seq != dbuf->win_base) ) {
			fprintf(stderr, "Oops : nothing to write, missing %u seq\n", dbuf->win_base);
			break;
		}

		/* Skip if next dgram is a dup */
		if ( ( vlen > 0 ) && (curr_seq == prev_seq) ) {
			continue;
		}

		/* Break if next seq dgram is missing */
		if ( ( vlen > 0 ) && (curr_seq > prev_seq+1 ) ) {
			break;
		}

		/* Normal case : curr_seq is the next dgram to write */
		dgram_index = dbuf->dgram_ordered_seq_numbers[i].index;
		dgram_len = dbuf->dgram_len[dgram_index] - dbuf->dgram_header_size;

		dbuf->write_iovecs[vlen].iov_len = dgram_len; /* Setup iovecs */
		dbuf->write_iovecs[vlen].iov_base = dbuf->buf + dgram_index*dbuf->dgram_max_size + dbuf->dgram_header_size;
		dbuf->dgram_seq_numbers[dgram_index] = 0; /* Mark dgram slots about to be written out as reusable */

		total += dgram_len; /* Update counters */
		vlen++;
		dbuf->win_base = curr_seq;
		prev_seq = curr_seq;
	}

	/* If nothing valid to write out */
	if ( vlen == 0 ) {
		return -1;
	}

	nwrite = writev(fd, dbuf->write_iovecs, vlen);
	if ( nwrite < 0 ) {
		perror("writev()");
		return nwrite;
	}

	if ( nwrite != total ) {
		fprintf(stderr, "writev() short\n");
		return nwrite;
	}

	return nwrite;
}

dgrambuf_t dgrambuf_new(size_t dgram_count, size_t dgram_max_size, size_t dgram_header_size) {

	dgrambuf_t dbuf = calloc(1, sizeof(struct dgrambuf_t));
	if (!dbuf) goto fail0;

	dbuf->dgram_count = dgram_count;
	dbuf->dgram_max_size = dgram_max_size;
	dbuf->dgram_header_size = dgram_header_size;

	dbuf->recv_iovecs = calloc(dgram_count, sizeof(struct iovec));
	if (!dbuf->recv_iovecs) goto fail1;

	dbuf->write_iovecs = calloc(dgram_count, sizeof(struct iovec));
	if (!dbuf->write_iovecs) goto fail2;

	dbuf->msgs = calloc(dgram_count, sizeof(struct mmsghdr));
	if (!dbuf->msgs) goto fail3;

	dbuf->win_base = 1;
	dbuf->dgram_seq_numbers = calloc(dgram_count, sizeof(unsigned int));
	if (!dbuf->dgram_seq_numbers) goto fail4;

	dbuf->dgram_len = calloc(dgram_count, sizeof(ssize_t));
	if (!dbuf->dgram_len) goto fail5;

	dbuf->dgram_ordered_seq_numbers = calloc(dgram_count, sizeof(struct uint_pair));
	if (!dbuf->dgram_ordered_seq_numbers) goto fail6;

	dbuf->buf = calloc(dgram_count, dgram_max_size);
	if (!dbuf->buf) goto fail7;

	return dbuf;

fail7:  free(dbuf->dgram_ordered_seq_numbers);
fail6:	free(dbuf->dgram_len);
fail5:	free(dbuf->dgram_seq_numbers);
fail4:	free(dbuf->msgs);
fail3:	free(dbuf->write_iovecs);
fail2:	free(dbuf->recv_iovecs);
fail1:	free(dbuf);
fail0:	return 0;
}

void dgrambuf_free(dgrambuf_t *dbuf) {
	if (dbuf && *dbuf) {
		free((*dbuf)->buf);
		free((*dbuf)->dgram_ordered_seq_numbers);
		free((*dbuf)->dgram_len);
		free((*dbuf)->dgram_seq_numbers);
		free((*dbuf)->msgs);
		free((*dbuf)->write_iovecs);
		free((*dbuf)->recv_iovecs);
		free(*dbuf);
	}
	*dbuf = NULL;
}

int _compare_uint_pair(const void *pa, const void *pb) {
	const struct uint_pair *a = pa;
	const struct uint_pair *b = pb;
	if (a->value < b->value)
		return -1;
	else if ( a->value > b->value )
		return 1;
	else
		return 0;
}