#include "common.h"
#include <math.h>
#include "e2e.h"

static uint16_t u16le(const char *src)
{
	uint8_t bytes[2];

	memcpy(bytes, src, 2);

	return (uint16_t)bytes[0]
	       + (((uint16_t)bytes[1]) << 8);
}

static uint32_t u32le(const char *src)
{
	uint8_t bytes[4];

	memcpy(bytes, src, 4);

	return (uint32_t)bytes[0]
	       + (((uint32_t)bytes[1]) << 8)
	       + (((uint32_t)bytes[2]) << 16)
	       + (((uint32_t)bytes[3]) << 24);
}

static int make_pointer(const struct e2e_data *data, const char *pos,
                        const char **out)
{
	uint32_t offs;

	offs = u32le(pos);

	if (!offs)
		*out = NULL;
	else {
		if (offs > data->end - data->start)
			return 1;

		*out = data->start + offs;
	}

	return 0;
}

#define HEADER_END 36

#define DIRECTORY_NUM_ENTRIES   36
#define DIRECTORY_CURRENT       40
#define DIRECTORY_NEXT          44
#define DIRECTORY_END          52

#define ENTRY_POS          0
#define ENTRY_START        4 
#define ENTRY_SIZE         8
#define ENTRY_PATIENT_ID   16
#define ENTRY_STUDY_ID     20
#define ENTRY_SERIES_ID    24
#define ENTRY_SLICE_ID     28
#define ENTRY_TYPE         36
#define ENTRY_END          44

#define CHUNK_POS          20
#define CHUNK_SIZE         24
#define CHUNK_PATIENT_ID   32
#define CHUNK_STUDY_ID     36
#define CHUNK_SERIES_ID    40
#define CHUNK_SLICE_ID     44
#define CHUNK_IND          48
#define CHUNK_TYPE         52
#define CHUNK_HEADER_END   60

#define IMAGE_SIZE      0
#define IMAGE_TYPE      4
#define IMAGE_WIDTH     12
#define IMAGE_HEIGHT    16
#define IMAGE_IMAGE     20

#define UF16_MANTISSA_MASK  ((1 << 11) - 1)  // 10 most significant bits
#define UF16_EXP_OFFSET -70 // FIXME: pure guess

static int read_fundus(struct e2e_data *data, const char *cursor,
                       struct e2e_image *image)
{
	if (cursor + image->width * image->height > data->end) {
		errorf("Fundus image at %td ends past the end of file.\n",
		        cursor - data->start);
		return 1;
	}

	// This cast is _technically_ illegal, but I can't think of a way in
	// which this would cause problems in practice.
	image->fundus = (const uint8_t*)cursor;

	return 0;
}

static int read_tomogram(struct e2e_data *data, const char *cursor,
                         struct e2e_image *image)
{
	if (cursor + image->width * image->height * 2 > data->end) {
		errorf("Tomogram image at %td ends past the end of file.\n",
		        cursor - data->start);
		return 1;
	}

	image->tomogram = calloc(image->width * image->height, sizeof(float));
	if (!image->tomogram)
		return ENOMEM;

	for (size_t y = 0; y < image->height; y++)
	for (size_t x = 0; x < image->width; x++) {
		size_t i = y * image->width + x;
		uint16_t raw;
		int exp, man;

		raw = u16le(cursor + i * 2);
		man = raw & UF16_MANTISSA_MASK;
		exp = (raw & ~UF16_MANTISSA_MASK) >> 10;
		image->tomogram[i] = ldexpf(man, exp2f(exp + UF16_EXP_OFFSET));
	}

	return 0;
}

static int read_image(struct e2e_data *data, const char *cursor,
                      struct e2e_image *image)
{
	if (cursor + IMAGE_IMAGE > data->end) {
		errorf("Image data at %td ends past the end of file.\n",
		        cursor - data->start);
		return 1;
	}

	image->size = u32le(cursor + IMAGE_SIZE);
	image->type = u32le(cursor + IMAGE_TYPE);
	image->width = u32le(cursor + IMAGE_WIDTH);
	image->height = u32le(cursor + IMAGE_HEIGHT);
	cursor += IMAGE_IMAGE;

	switch (image->type) {
	case E2E_IMAGE_FUNDUS:
		return read_fundus(data, cursor, image);

	case E2E_IMAGE_TOMOGRAM:
		return read_tomogram(data, cursor, image);
	}

	return 0;
}

static int read_entry(struct e2e_data *data, const char *cursor,
                      struct e2e_entry *entry)
{
	const char *pos;
	struct e2e_chunk *chunk;

	if (cursor + ENTRY_END > data->end) {
		errorf("Entry at %td data->ends past the data->end of file.\n",
		        cursor - data->start);
		return 1;
	}

	if (make_pointer(data, cursor + ENTRY_POS, &pos)) {
		errorf("Entry at %td has a bad 'pos' entry.\n",
		        cursor - data->start);
		return 1;
	}

	if (pos != cursor) {
		errorf("Entry at %td has a 'pos' entry that differs from the entry's actual position",
		        cursor - data->start);
		return 1;
	}

	if (make_pointer(data, cursor + ENTRY_START, &entry->start)) {
		errorf("Entry at %td has a bad 'data->start' entry.\n",
		        cursor - data->start);
		return 1;
	}

	entry->size = u32le(cursor + ENTRY_SIZE);
	entry->patient_id = u32le(cursor + ENTRY_PATIENT_ID);
	entry->study_id = u32le(cursor + ENTRY_STUDY_ID);
	entry->series_id = u32le(cursor + ENTRY_SERIES_ID);
	entry->slice_id = u32le(cursor + ENTRY_SLICE_ID);
	entry->type = u32le(cursor + ENTRY_TYPE);

	if (!entry->start || entry->start < data->start + HEADER_END)
		return 0;

	cursor = entry->start;
	chunk = &entry->chunk;
	entry->has_chunk = 1;

	chunk->entry = entry;
	chunk->origin = cursor;

	if (cursor + CHUNK_HEADER_END > data->end) {
		errorf("Chunk at %td ends past the end of file.\n",
		        cursor - data->start);
		return 1;
	}

	if (make_pointer(data, cursor + CHUNK_POS, &pos)) {
		errorf("Chunk at %td has a bad 'pos' entry.\n",
		        cursor - data->start);
		return 1;
	}

	if (pos != cursor) {
		errorf("Chunk at %td has a 'pos' entry that differs from the entry's actual position.\n",
		        cursor - data->start);
		return 1;
	}

	chunk->size = u32le(cursor + CHUNK_SIZE);
	chunk->patient_id = u32le(cursor + CHUNK_PATIENT_ID);
	chunk->study_id = u32le(cursor + CHUNK_STUDY_ID);
	chunk->series_id = u32le(cursor + CHUNK_SERIES_ID);
	chunk->slice_id = u32le(cursor + CHUNK_SLICE_ID);
	chunk->type = u32le(cursor + CHUNK_TYPE);
	cursor += CHUNK_HEADER_END;

	switch (chunk->type) {
	case E2E_CHUNK_IMAGE:
		return read_image(data, cursor, &chunk->image);
	}

	return 0;
}

static int read_directory(struct e2e_data *data, const char *cursor,
                          struct e2e_directory *dir)
{
	if (cursor + DIRECTORY_END > data->end) {
		errorf("Directory at %td data->ends past the data->end of file.\n",
		        cursor - data->start);
		return 1;
	}

	dir->num_entries = u32le(cursor + DIRECTORY_NUM_ENTRIES);

	if (make_pointer(data, cursor + DIRECTORY_CURRENT, &dir->current))  {
		errorf("Directory at %td contains a bad 'current' field.\n",
		        cursor - data->start);
		return 1;
	}

	if (make_pointer(data, cursor + DIRECTORY_NEXT, &dir->next))  {
		errorf("Directory at %td contains a bad 'current' field.\n",
		        cursor - data->start);
		return 1;
	}

	return 0;
}

// Because entries are a part of a larger array in e2e_data, they're freed
// all at once by e2e_destroy and not here.
void e2e_entry_destroy(struct e2e_entry *entry)
{
	if (entry->has_chunk
	    && entry->chunk.type == E2E_CHUNK_IMAGE
	    && entry->chunk.image.type == E2E_IMAGE_TOMOGRAM)
		free(entry->chunk.image.tomogram);
}

void e2e_destroy(struct e2e_data *data)
{
	struct e2e_directory *dir, *next;
	struct e2e_patient *patient, *pnext;

	for (dir = data->dirs; dir; dir = next) {
		next = dir->data_list.next;

		for (size_t i = 0; i < dir->num_entries; i++)
			e2e_entry_destroy(dir->entries + i);
		free(dir->entries);
		free(dir);
	}

	HASH_ITER(hh, data->patients, patient, pnext) {
		struct e2e_study *study, *stnext;

		HASH_ITER(hh, patient->studies, study, stnext) {
			struct e2e_series *series, *ssnext;

			HASH_ITER(hh, study->series, series, ssnext) {
				struct e2e_slice *slice, *slnext;

				HASH_ITER(hh, series->slices, slice, slnext) {
					HASH_DEL(series->slices, slice);
					free(slice);
				}

				HASH_DEL(study->series, series);
				free(series);
			}

			HASH_DEL(patient->studies, study);
			free(study);
		}

		HASH_DEL(data->patients, patient);
		free(patient);
	}
}

static int ensure_slice(struct e2e_data *data, int patient_id, int study_id,
                         int series_id, int slice_id, struct e2e_patient **o_patient,
                         struct e2e_study **o_study, struct e2e_series **o_series,
                         struct e2e_slice **o_slice)
{
	struct e2e_patient *patient = NULL;
	_Bool new_patient = 0;
	struct e2e_study *study = NULL;
	_Bool new_study = 0;
	struct e2e_series *series = NULL;
	_Bool new_series = 0;
	struct e2e_slice *slice = NULL;

	if (patient_id == E2E_NONE)
		goto out;

	HASH_FIND_INT(data->patients, &patient_id, patient);
	if (!patient) {
		new_patient = 1;

		patient = calloc(1, sizeof(struct e2e_patient));
		if (!patient)
			goto error_patient;
		patient->id = patient_id;

		HASH_ADD_INT(data->patients, id, patient);
	}

	if (study_id == E2E_NONE)
		goto out;

	if (!new_patient)
		HASH_FIND_INT(patient->studies, &study_id, study);
	if (!study) {
		new_study = 1;

		study = calloc(1, sizeof(struct e2e_study));
		if (!study)
			goto error_study;
		study->id = study_id;
		HASH_ADD_INT(patient->studies, id, study);
	}

	if (series_id == E2E_NONE)
		goto out;

	if (!new_study)
		HASH_FIND_INT(study->series, &series_id, series);
	if (!series) {
		new_series = 1;

		series = calloc(1, sizeof(struct e2e_series));
		if (!series)
			goto error_series;
		series->id = series_id;

		HASH_ADD_INT(study->series, id, series);
	}

	if (slice_id == E2E_NONE)
		goto out;

	if (!new_series)
		HASH_FIND_INT(series->slices, &slice_id, slice);
	if (!slice) {
		slice = calloc(1, sizeof(struct e2e_slice));
		if (!slice)
			goto error_slice;
		slice->id = slice_id;

		HASH_ADD_INT(series->slices, id, slice);
	}

out:
	if (o_patient)
		*o_patient = patient;
	if (o_study)
		*o_study = study;
	if (o_series)
		*o_series = series;
	if (o_slice)
		*o_slice = slice;

	return 0;

error_slice:
	if (new_series)
		free(series);
error_series:
	if (new_study)
		free(study);
error_study:
	if (new_patient)
		free(patient);
error_patient:
	return ENOMEM;
}

// For now the file is loaded as-is, ie. following the original 'structure'.
// This function then untangles the mess so we get something usable to work with.
static int post_process(struct e2e_data *data)
{
	int rv;
	struct e2e_directory *dir;

	eli_for(dir, data->dirs, data_list)
	for (size_t i = 0; i < dir->num_entries; i++) {
		struct e2e_entry *entry = dir->entries + i;
		struct e2e_series *series;
		struct e2e_slice *slice;

		rv = ensure_slice(data, entry->patient_id, entry->study_id,
		                   entry->series_id, entry->slice_id,
		                   NULL, NULL, &series, &slice);
		if (rv)
			return rv;

		if (entry->has_chunk && entry->chunk.type == E2E_CHUNK_IMAGE) {
			struct e2e_image *image = &entry->chunk.image;

			if (slice &&
			    entry->chunk.image.type == E2E_IMAGE_TOMOGRAM)
				slice->image = image;

			if (series &&
			    entry->chunk.image.type == E2E_IMAGE_FUNDUS) {
				eli_append(&series->fundi, image,
				           series_fundi_list);
				series->num_fundi++;
			}
		}
	}

	return 0;
}

int e2e_read(struct e2e_data *data, const char *start, const char *end)
{
	int rv = 0;
	const char *cursor;

	data->start = start;
	data->end = end;
	data->dirs = NULL;
	data->patients = NULL;

	if (start + HEADER_END > end) {
		errorf("File too short to be an E2E file.\n");
		return 1;
	}

	// TODO: verify the header

	cursor = start + HEADER_END;
	cursor += DIRECTORY_END; // skip the first directory (wtf?)

	while (cursor) {
		struct e2e_directory *dir;
		size_t i;

		dir = malloc(sizeof(struct e2e_directory));
		if (!dir) {
			rv = ENOMEM;
			goto error;
		}

		rv = read_directory(data, cursor, dir);
		if (rv)
			goto error_read_directory;

		dir->entries = calloc(dir->num_entries, sizeof(struct e2e_entry));
		if (!dir->entries) {
			rv = ENOMEM;
			goto error_alloc_entries;
		}

		for (i = 0; i < dir->num_entries; i++) {
			const char *entry_cursor;
			struct e2e_entry *entry = dir->entries + i;

			entry_cursor = cursor + DIRECTORY_END + i * ENTRY_END;
			rv = read_entry(data, entry_cursor, entry);
			if (rv)
				goto error_read_entry;
		}

		eli_append(&data->dirs, dir, data_list);
		cursor = dir->next;
		continue;


	error_read_entry:
		for (size_t j = 0; j < i; j++)
			e2e_entry_destroy(dir->entries + j);
		free(dir->entries);
	error_alloc_entries:
	error_read_directory:
		free(dir);
		goto error;
	}

	rv = post_process(data);
	if (rv)
		goto error;

	return 0;

error:
	e2e_destroy(data);
	return rv;
}