#include "model_format_utah.h"

#include <assert.h>
#include <ctype.h>
#include <fcntl.h>
#include <math.h>
#include <stdbit.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#define BUFFER_SIZE 256

#define FIB_WRAP_COEFF 11400714819323198485LLU

typedef struct {
	uint8_t *rp;
	size_t vertex_count;
	size_t index_count;
	size_t capacity;
	size_t buffered;
	const int fd;
	uint8_t buffer[BUFFER_SIZE];
} read_ctx_t;

typedef struct {
	unsigned probe_count;
	index_t index;
	bool occupied;
} table_node_t;

typedef struct {
	table_node_t *nodes;
	size_t size;
	unsigned shift;
	unsigned max_probe_count;
} table_t;

static uint32_t hash_vertex(vertex_t v)
{
	const uint32_t px = *((uint32_t *)&v.position.x);
	const uint32_t py = *((uint32_t *)&v.position.y);
	const uint32_t pz = *((uint32_t *)&v.position.z);
	const uint32_t nx = *((uint32_t *)&v.normal.x);
	const uint32_t ny = *((uint32_t *)&v.normal.y);
	const uint32_t nz = *((uint32_t *)&v.normal.z);
	return px ^ (py << 1) ^ (pz << 2) ^ (pz << 3) ^ (nx << 4) ^ (ny << 5)
	    ^ (nz << 6);
}

static bool vertex_equals(vertex_t v, vertex_t u)
{
	return v.position.x == u.position.x && v.position.y == u.position.y
	    && v.position.z == u.position.z && v.normal.x == u.normal.x
	    && v.normal.y == u.normal.y && v.normal.z == u.normal.z;
}

static unsigned wrap(table_t *table, uint32_t hash)
{
	return (unsigned)((hash * FIB_WRAP_COEFF) >> table->shift);
}

static bool lookup(
    table_t *table, model_t *model, unsigned base, vertex_t vertex,
    index_t *index_out)
{
	for (unsigned i = 0; i < table->max_probe_count; ++i) {
		const unsigned t = (base + i) & (UINT64_MAX >> table->shift);
		const table_node_t *node = table->nodes + t;
		if (!node->occupied)
			return false;

		const vertex_t node_vertex = model->vertices[node->index];
		if (vertex_equals(vertex, node_vertex)) {
			*index_out = node->index;
			return true;
		}
	}
	return false;
}

static void insert(table_t *table, unsigned base, index_t index)
{
	for (unsigned i = 0; i < table->size; ++i) {
		const unsigned t = (base + i) & (UINT64_MAX >> table->shift);
		table_node_t *node = table->nodes + t;

		if (!node->occupied) {
			node->occupied = true;
			node->probe_count = i;
			node->index = index;
			if (i > table->max_probe_count)
				table->max_probe_count = i;
			return;
		}

		if (node->probe_count < i) {
			const unsigned node_probe_count = node->probe_count;
			node->probe_count = i;
			i = node_probe_count;
			const index_t node_index = node->index;
			node->index = index;
			index = node_index;
		}
	}

	assert(false);
}

static index_t handle_vertex(table_t *table, model_t *model, vertex_t vertex)
{
	index_t index;
	const unsigned base = wrap(table, hash_vertex(vertex));
	const bool exists = lookup(table, model, base, vertex, &index);
	if (exists)
		return index;
	index = model->vertex_count;
	model->vertices[model->vertex_count++] = vertex;
	insert(table, base, index);
	return index;
}

static uint8_t read_byte(read_ctx_t *ctx)
{
	if (ctx->rp >= ctx->buffer + ctx->buffered) {
		const ssize_t res = read(ctx->fd, ctx->buffer, BUFFER_SIZE);
		assert(res != -1);
		ctx->buffered = (size_t)res;
		ctx->rp = ctx->buffer;
	}
	return *ctx->rp++;
}

static size_t read_triangle_count(read_ctx_t *ctx)
{
	size_t triangle_count = 0;
	while (1) {
		const uint8_t b = read_byte(ctx);
		if (isdigit(b)) {
			triangle_count = 10 * triangle_count + b - '0';
		} else {
			assert(b == '\n');
			break;
		}
	}
	return triangle_count;
}

static float read_float(read_ctx_t *ctx)
{
	bool negative = false;
	float value = 0, digit_div = 0;
	while (1) {
		const uint8_t b = read_byte(ctx);
		if (b == '-') {
			assert(value == 0);
			negative = true;
		} else if (b == '.') {
			digit_div = 10;
		} else if (isdigit(b)) {
			if (digit_div == 0) {
				value = 10 * value + (b - '0');
			} else {
				value += (float)(b - '0') / digit_div;
				digit_div *= 10;
			}
		} else {
			assert(b == ' ' || b == '\n');
			break;
		}
	}
	return (negative ? -1 : 1) * value;
}

static vec3_t read_vec3(read_ctx_t *ctx)
{
	const float x = read_float(ctx);
	const float y = read_float(ctx);
	const float z = read_float(ctx);
	return (vec3_t) { x, y, z };
}

static vertex_t read_vertex(read_ctx_t *ctx)
{
	const vec3_t position = read_vec3(ctx);
	const vec3_t normal = read_vec3(ctx);
	return (vertex_t) { position, normal, {}, {} };
}

model_t model_format_utah_read(const char *path)
{
	read_ctx_t ctx = { .fd = open(path, O_RDONLY), .rp = ctx.buffer };
	assert(ctx.fd != -1);

	const size_t triangle_count = read_triangle_count(&ctx);
	const size_t index_count = 3 * triangle_count;

	// Begin with vertex buffer at maximum possible size
	model_t model = model_init(index_count, index_count);
	model.vertex_count = 0;

	// Set up hash table for vertex de-duplication
	table_t table = { .max_probe_count = 0 };
	table.size = stdc_bit_ceil((size_t)(1.5 * index_count));
	table.shift = 64 - stdc_trailing_zeros(table.size);
	table.nodes = calloc(table.size, sizeof(table_node_t));
	assert(table.nodes != nullptr);

	for (size_t i = 0; i < index_count; ++i) {
		const vertex_t vertex = read_vertex(&ctx);
		index_t index = handle_vertex(&table, &model, vertex);
		model.indices[i] = index;

		// Skip newline after each triangle
		if (i % 3 == 2) {
			const uint8_t b = read_byte(&ctx);
			assert(b == '\n');
		}
	}

	// Shrink vertex buffer down to needed size
	model.vertices
		= realloc(model.vertices, model.vertex_count * sizeof(vertex_t));
	assert(model.vertices != nullptr);

	free(table.nodes);
	close(ctx.fd);
	return model;
}