#define _DEFAULT_SOURCE

#include "idx.h"

#include <endian.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdint.h>
#include <stdarg.h>

static const uint8_t idx_type_encode_lut[IDX_TYPE_COUNT] = {
	[IDX_TYPE_UNSIGNED_BYTE] = 0x08,
	[IDX_TYPE_SIGNED_BYTE] = 0x09,
	[IDX_TYPE_SHORT] = 0x0B,
	[IDX_TYPE_INT] = 0x0C,
	[IDX_TYPE_FLOAT] = 0x0D,
	[IDX_TYPE_DOUBLE] = 0x0E,
};

static const uint8_t idx_type_size[IDX_TYPE_COUNT] = {
	[IDX_TYPE_UNSIGNED_BYTE] = 1,
	[IDX_TYPE_SIGNED_BYTE] = 1,
	[IDX_TYPE_SHORT] = 2,
	[IDX_TYPE_INT] = 4,
	[IDX_TYPE_FLOAT] = 4,
	[IDX_TYPE_DOUBLE] = 8
};

static bool valid_idx_type(int type);

static int idx_type_decode(int val);

static bool fread_u32(FILE *f, uint32_t *val);
static bool fwrite_u32(FILE *f, uint32_t val);

bool
valid_idx_type(int type)
{
	return type >= 0 && type < IDX_TYPE_COUNT;
}

int
idx_type_decode(int val)
{
	int i;

	for (i = 0; i < IDX_TYPE_COUNT; i++) {
		if (val == idx_type_encode_lut[i])
			return i;
	}

	return IDX_TYPE_UNKNOWN;
}

bool
fread_u32(FILE *f, uint32_t *val)
{
	size_t len;

	len = fread(val, 4, 1, f);
	if (len != 1) return false;
	*val = be32toh(*val);

	return true;
}

bool
fwrite_u32(FILE *f, uint32_t val)
{
	size_t len;

	val = htobe32(val);
	len = fwrite(&val, 4, 1, f);
	if (len != 1) return false;

	return true;
}

int
idx_new(struct idx *idx, uint8_t dims_count, uint32_t *dims, int data_type,
	const struct allocator *allocator)
{
	size_t i, size;
	int rc;

	if (!valid_idx_type(data_type))
		return IDX_BAD_PARAM;

	size = idx_type_size[data_type];
	for (i = 0; i < dims_count; i++)
		size *= dims[i];

	idx->allocator = allocator;
	idx->buf.data = allocator->alloc(allocator, size, &rc);
	if (!idx->buf.data) return -rc;
	idx->buf.len = size;

	idx->dims.sizes = dims; /* move */
	idx->dims.count = dims_count;

	idx->type.size = idx_type_size[data_type];
	idx->type.val = data_type;

	return IDX_OK;
}

int
idx_wrap(struct idx *idx, void *data, uint8_t dims_count,
	uint32_t *dims, int data_type)
{
	size_t i, size;

	if (!valid_idx_type(data_type))
		return IDX_BAD_PARAM;

	size = idx_type_size[data_type];
	for (i = 0; i < dims_count; i++)
		size *= dims[i];

	idx->allocator = NULL;

	idx->buf.data = data; /* move */
	idx->buf.len = size;

	idx->dims.sizes = dims; /* move */
	idx->dims.count = dims_count;

	idx->type.size = idx_type_size[data_type];
	idx->type.val = data_type;

	return IDX_OK;
}

int
idx_load(struct idx *idx, const char *filename,
	const struct allocator *allocator)
{
	uint32_t magic;
	size_t len, i;
	FILE *file;
	int rc;

	file = fopen(filename, "r");
	if (file == NULL)
		return IDX_FILE_ACCESS;

	/* read magic */
	if (!fread_u32(file, &magic))
		return IDX_FILE_READ_MAGIC;

	if ((magic & 0xFFFF0000) != 0)
		return IDX_FILE_BAD_MAGIC;

	idx->type.val = idx_type_decode((magic & 0xFF00) >> 8);
	if (idx->type.val == IDX_TYPE_UNKNOWN)
		return IDX_FILE_BAD_TYPE;
	idx->type.size = idx_type_size[idx->type.val];

	idx->dims.count = magic & 0xFF;
	if (idx->dims.count == 0)
		return IDX_FILE_BAD_DIMS;

	idx->allocator = allocator;
	idx->dims.sizes = allocator->alloc(allocator,
		sizeof(uint32_t) * idx->dims.count, &rc);
	if (!idx->dims.sizes) return -rc;

	/* read dims */
	idx->buf.len = 1;
	for (i = 0; i < idx->dims.count; i++) {
		if (!fread_u32(file, &idx->dims.sizes[i]))
			return IDX_FILE_READ_DIMS;
		idx->buf.len *= idx->dims.sizes[i];
	}

	idx->buf.data = allocator->alloc(allocator, idx->buf.len, &rc);
	if (idx->buf.data) return -rc;

	/* read data */
	len = fread(idx->buf.data, idx->type.size, idx->buf.len, file);
	if (len != idx->buf.len)
		return IDX_FILE_READ_DATA;

	fclose(file);

	return IDX_OK;
}

int
idx_save(struct idx *idx, const char *filename)
{
	FILE *file;
	uint32_t val;
	size_t i, len, size;

	file = fopen(filename, "w+");
	if (file == NULL)
		return IDX_FILE_ACCESS;

	if (!valid_idx_type(idx->type.val))
		return IDX_BAD_PARAM;
	val = idx_type_encode_lut[idx->type.val];

	/* write magic */
	if (!fwrite_u32(file, (val << 8) + idx->dims.count))
		return IDX_FILE_WRITE_MAGIC;

	/* write dims */
	for (i = 0; i < idx->dims.count; i++) {
		if (!fwrite_u32(file, idx->dims.sizes[i]))
			return IDX_FILE_WRITE_DIMS;
	}

	size = idx_type_size[idx->type.val];
	for (i = 0; i < idx->dims.count; i++)
		size *= idx->dims.sizes[i];

	/* write data */
	len = fwrite(idx->buf.data, 1, size, file);
	if (len != size) return IDX_FILE_WRITE_DATA;

	fclose(file);

	return IDX_OK;
}

int
idx_deinit(struct idx *idx)
{
	int rc;

	if (!idx->allocator) return IDX_OK;

	rc = idx->allocator->free(idx->allocator, idx->buf.data);
	if (rc) return -rc;

	rc = idx->allocator->free(idx->allocator, idx->dims.sizes);
	if (rc) return -rc;

	return IDX_OK;
}

void *
idx_get(struct idx *idx, int dims, ...)
{
	va_list ap;
	size_t stride;
	size_t i, val;
	void *p;

	if (dims < 0 || dims > idx->dims.count)
		return NULL;

	stride = 1;
	for (i = 0; i < idx->dims.count; i++)
		stride *= idx->dims.sizes[i];

	p = idx->buf.data;
	va_start(ap, dims);
	for (i = 0; i < dims; i++) {
		val = va_arg(ap, size_t);
		if (val >= idx->dims.sizes[i])
			return NULL;

		stride /= idx->dims.sizes[i];
		p += stride * val;
	}
	va_end(ap);

	return p;
}