#include <string.h>

#include "preamble.h"
#include "config.h"

#include "packet_mod.h"

result_t packet_mod_init(packet_mod_ctx_t *ctx)
{
	ctx->pkt_bytes = NULL;
	ctx->pkt_symbols = NULL;

	ctx->length = 0;

	ctx->fec   = fec_create(CHANNEL_CODE, NULL);
	ctx->modem = modem_create(MODULATION);

	ctx->hdr_fec   = fec_create(HEADER_CHANNEL_CODE, NULL);
	ctx->hdr_modem = modem_create(HEADER_MODULATION);

	ctx->state = NOT_STARTED;

	return OK;
}


result_t packet_mod_free(packet_mod_ctx_t *ctx)
{
	switch(ctx->state) {
		case DATA_RAW:
		case DATA_CODED:
			free(ctx->pkt_bytes);
			ctx->pkt_bytes = NULL;
			break;

		case DATA_MODULATED:
		case HEADER_ADDED:
		case PREAMBLE_ADDED:
			free(ctx->pkt_symbols);
			ctx->pkt_symbols = NULL;
			break;

		case NOT_STARTED:
			// nothing to do
			break;
	}

	ctx->state = NOT_STARTED;
	ctx->length = 0;

	fec_destroy(ctx->fec);
	modem_destroy(ctx->modem);

	fec_destroy(ctx->hdr_fec);
	modem_destroy(ctx->hdr_modem);

	return OK;
}


result_t packet_mod_set_data(packet_mod_ctx_t *ctx, const unsigned char *data, size_t length)
{
	if(ctx->state != NOT_STARTED) {
		return ERR_INVALID_STATE;
	}

	ctx->pkt_bytes = malloc(length);
	if(!ctx->pkt_bytes) {
		return ERR_NO_MEM;
	}

	memcpy(ctx->pkt_bytes, data, length);
	ctx->length = length;

	ctx->state = DATA_RAW;

	return OK;
}


result_t packet_mod_encode(packet_mod_ctx_t *ctx)
{
	if(ctx->state != DATA_RAW) {
		return ERR_INVALID_STATE;
	}

	unsigned int enc_length = fec_get_enc_msg_length(CHANNEL_CODE, ctx->length);

	unsigned char *enc_msg = malloc(enc_length);
	if(!enc_msg) {
		return ERR_NO_MEM;
	}

	ERR_CHECK_LIQUID(fec_encode(ctx->fec, ctx->length, ctx->pkt_bytes, enc_msg));

	free(ctx->pkt_bytes);

	ctx->pkt_bytes = enc_msg;
	ctx->length = enc_length;

	ctx->state = DATA_CODED;

	return OK;
}


result_t packet_mod_modulate(packet_mod_ctx_t *ctx)
{
	if(!((ctx->state == DATA_RAW) || (ctx->state == DATA_CODED))) {
		return ERR_INVALID_STATE;
	}

	unsigned int bps = modem_get_bps(ctx->modem);
	unsigned int nsyms = (ctx->length * 8 + bps/2) / bps;

	unsigned char msg_sym_idcs[nsyms];

	float complex *msg_mod = malloc(sizeof(float complex) * nsyms);
	if(!msg_mod) {
		return ERR_NO_MEM;
	}

	ERR_CHECK_LIQUID(liquid_repack_bytes(ctx->pkt_bytes, 8, ctx->length, msg_sym_idcs, bps, nsyms, &nsyms));

	for(size_t i = 0; i < nsyms; i++) {
		ERR_CHECK_LIQUID(modem_modulate(ctx->modem, msg_sym_idcs[i], &msg_mod[i]));
	}

	free(ctx->pkt_bytes);
	ctx->pkt_bytes = NULL;

	ctx->pkt_symbols = msg_mod;
	ctx->length = nsyms;

	ctx->state = DATA_MODULATED;

	return OK;
}


result_t packet_mod_add_header(packet_mod_ctx_t *ctx)
{
	if(ctx->state != DATA_MODULATED) {
		return ERR_INVALID_STATE;
	}

	if(ctx->length >= (1 << 16)) {
		// packet length cannot be stored in the uint16_t in the header.
		return ERR_SIZE;
	}

	// build the header. All field are transferred in network byte order (big endian).
	uint8_t header[4];

	// set length
	header[0] = (ctx->length >> 8) & 0xFF;
	header[1] = (ctx->length >> 0) & 0xFF;
	// set raw data CRC
	header[2] = (ctx->raw_data_crc >> 8) & 0xFF;
	header[3] = (ctx->raw_data_crc >> 0) & 0xFF;

	// note: the header is coded and modulated differently than the data.

	// encode the header
	unsigned int hdr_enc_length = fec_get_enc_msg_length(HEADER_CHANNEL_CODE, sizeof(header));
	unsigned char header_encoded[hdr_enc_length];

	ERR_CHECK_LIQUID(fec_encode(ctx->hdr_fec, sizeof(header), header, header_encoded));

	// modulate the header
	unsigned int bps = modem_get_bps(ctx->hdr_modem);
	unsigned int nsyms = (hdr_enc_length * 8 + bps/2) / bps;

	unsigned char header_sym_idcs[nsyms];
	float complex header_mod[nsyms];

	ERR_CHECK_LIQUID(liquid_repack_bytes(header_encoded, 8, hdr_enc_length, header_sym_idcs, bps, nsyms, &nsyms));

	for(size_t i = 0; i < nsyms; i++) {
		ERR_CHECK_LIQUID(modem_modulate(ctx->hdr_modem, header_sym_idcs[i], &header_mod[i]));
	}

	// concatenate the modulated header and message
	size_t packet_length_with_header = ctx->length + nsyms;
	float complex *packet_with_header = malloc(sizeof(float complex) * packet_length_with_header);
	if(!packet_with_header) {
		return ERR_NO_MEM;
	}

	memcpy(packet_with_header, header_mod, sizeof(float complex) * nsyms);
	memcpy(packet_with_header+nsyms, ctx->pkt_symbols, sizeof(float complex) * ctx->length);

	// replace the data in the context struct
	free(ctx->pkt_symbols);

	ctx->pkt_symbols = packet_with_header;
	ctx->length = packet_length_with_header;

	ctx->state = HEADER_ADDED;

	return OK;
}


result_t packet_mod_add_preamble(packet_mod_ctx_t *ctx)
{
	if(ctx->state != HEADER_ADDED) {
		return ERR_INVALID_STATE;
	}

	// concatenate the preamble with the message
	size_t pre_length = preamble_get_symbol_count();

	size_t packet_length_with_preamble = ctx->length + pre_length;
	float complex *packet_with_preamble = malloc(sizeof(float complex) * packet_length_with_preamble);
	if(!packet_with_preamble) {
		return ERR_NO_MEM;
	}

	memcpy(packet_with_preamble, preamble_get_symbols(), sizeof(float complex) * pre_length);
	memcpy(packet_with_preamble+pre_length, ctx->pkt_symbols, sizeof(float complex) * ctx->length);

	// replace the data in the context struct
	free(ctx->pkt_symbols);

	ctx->pkt_symbols = packet_with_preamble;
	ctx->length = packet_length_with_preamble;

	ctx->state = PREAMBLE_ADDED;

	return OK;
}


result_t packet_mod_get_result_b(packet_mod_ctx_t *ctx, unsigned char *data, size_t *length)
{
	if(!((ctx->state == DATA_RAW) || (ctx->state == DATA_CODED))) {
		return ERR_INVALID_STATE;
	}

	if(!data) {
		// data is NULL, so we only return the required length
		*length = ctx->length;
		return OK;
	}

	if(*length < ctx->length) {
		return ERR_NO_MEM;
	}

	*length = ctx->length;
	memcpy(data, ctx->pkt_bytes, *length * sizeof(unsigned char));

	return OK;
}


result_t packet_mod_get_result_cf(packet_mod_ctx_t *ctx, float complex *data, size_t *length)
{
	if(!((ctx->state == DATA_MODULATED) || (ctx->state == HEADER_ADDED) || (ctx->state == PREAMBLE_ADDED))) {
		return ERR_INVALID_STATE;
	}

	if(!data) {
		// data is NULL, so we only return the required length
		*length = ctx->length;
		return OK;
	}

	if(*length < ctx->length) {
		return ERR_NO_MEM;
	}

	*length = ctx->length;
	memcpy(data, ctx->pkt_symbols, *length * sizeof(float complex));

	return OK;
}