#include "model_format_glb.h"

#include <assert.h>
#include <ctype.h>
#include <fcntl.h>
#include <stdio.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

/*
 * Ensure the system is little-endian.
 *
 * GLB is a little-endian format and this module assumes the host
 * system matches this in multiple places.
 */
#if defined(__BYTE_ORDER__) && defined(__ORDER_LITTLE_ENDIAN__)
static_assert(
    __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__,
    "This module requires a little-endian system");
#else
#warning "Could not determine endianness: assuming little endian"
#endif

#define MAGIC 0x46546c67 // "glTF"
#define VERSION 2

#define MAX_KEY_SIZE 128

enum {
	CHUNK_TYPE_JSON = 0x4e4f534a, // "JSON"
	CHUNK_TYPE_BIN = 0x004e4942, // "BIN\0"
};

enum {
	RENDERING_MODE_TRIANGLES = 4,
};

typedef struct {
	uint32_t magic;
	uint32_t version;
	uint32_t len;
} file_hdr_t;

typedef struct {
	uint32_t len;
	uint32_t type;
} chunk_hdr_t;

typedef struct {
	uint32_t normals;
	uint32_t positions;
	uint32_t indices;
} mesh_t;

typedef struct {
	uint32_t view;
	uint32_t offset;
	uint32_t count;
} accessor_t;

typedef struct {
	accessor_t normals;
	accessor_t positions;
	accessor_t indices;
} accessors_t;

typedef struct {
	uint32_t offset;
	uint32_t len;
	uint32_t stride;
} view_t;

typedef struct {
	view_t normals;
	view_t positions;
	view_t indices;
} views_t;

typedef struct {
	uint32_t offset;
	uint32_t len;
} gltf_info_member_t;

typedef struct {
	gltf_info_member_t normals;
	gltf_info_member_t positions;
	gltf_info_member_t indices;
} gltf_info_t;

static void *map_file(const char *path)
{
	const int fd = open(path, O_RDONLY);
	assert(fd != -1);

	struct stat stat;
	const int res = fstat(fd, &stat);
	assert(res == 0);

	void *addr = mmap(nullptr, stat.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
	assert(addr != nullptr);

	close(fd);
	return addr;
}

static int read_key(const uint8_t **p, char buffer[MAX_KEY_SIZE])
{
	assert(**p == '"');
	++(*p);
	size_t len = 0;
	while (**p != '"') {
		assert(len < MAX_KEY_SIZE);
		buffer[len++] = *(*p)++;
	}
	++(*p);
	return len;
}

static bool ends_value(uint8_t b)
{
	switch (b) {
	case '}':
	case ']':
	case ',':
		return true;
	default:
		return false;
	}
}

static void skip_value(const uint8_t **p)
{
	switch (**p) {
		int nesting;
	case '{':
		++(*p);
		nesting = 0;
		while (**p != '}' || nesting > 0) {
			if (**p == '{')
				++nesting;
			if (**p == '}')
				--nesting;
			++(*p);
		}
		++(*p);
		break;
	case '[':
		++(*p);
		nesting = 0;
		while (**p != ']' || nesting > 0) {
			if (**p == '[')
				++nesting;
			if (**p == ']')
				--nesting;
			++(*p);
		}
		++(*p);
		break;
	case '"':
		++(*p);
		while (**p != '"') {
			if (**p == '\\')
				++(*p);
			++(*p);
		}
		++(*p);
		break;
	default:
		++(*p);
		while (!ends_value(**p))
			++(*p);
		break;
	}
}

static uint32_t read_uint(const uint8_t **p)
{
	uint32_t value = 0;
	while (!ends_value(**p)) {
		assert(isdigit(**p));
		value = 10 * value + *(*p)++ - '0';
	}
	return value;
}

static mesh_t read_mesh_primitives(const uint8_t **p)
{
	char buffer[MAX_KEY_SIZE];
	size_t len;
	mesh_t mesh = {};

	assert(**p == '[');
	++(*p);

	assert(**p == '{');
	++(*p);

	while (**p != '}') {
		len = read_key(p, buffer);
		assert(**p == ':');
		++(*p);

		if (memcmp("attributes", buffer, len) == 0) {
			assert(**p == '{');
			++(*p);

			while (**p != '}') {
				len = read_key(p, buffer);
				assert(**p == ':');
				++(*p);
				if (memcmp("NORMAL", buffer, len) == 0) {
					mesh.normals = read_uint(p);
				} else if (memcmp("POSITION", buffer, len) == 0) {
					mesh.positions = read_uint(p);
				} else {
					skip_value(p);
				}
				if (**p == ',')
					++(*p);
			}
			++(*p);
		} else if (memcmp("indices", buffer, len) == 0) {
			mesh.indices = read_uint(p);
		} else if (memcmp("mode", buffer, len) == 0) {
			assert(read_uint(p) == RENDERING_MODE_TRIANGLES);
		} else {
			skip_value(p);
		}

		if (**p == ',')
			++(*p);
	}
	++(*p);

	assert(**p == ']');
	++(*p);

	return mesh;
}

static mesh_t read_meshes(const uint8_t *p)
{
	char buffer[MAX_KEY_SIZE];
	size_t len;

	assert(*p == '[');
	++p;
	assert(*p == '{');
	++p;

	while (*p != '}') {
		len = read_key(&p, buffer);
		assert(*p == ':');
		++p;
		if (memcmp("primitives", buffer, len) == 0)
			return read_mesh_primitives(&p);
		else
			skip_value(&p);
		if (*p == ',')
			++p;
	}

	/*
	 * If we reach here, the primitives key was missing from the mesh
	 * definition.
	 */
	assert(false);
}

static accessor_t read_accessor(const uint8_t **p)
{
	char buffer[MAX_KEY_SIZE];
	size_t len;
	accessor_t accessor = {};

	assert(**p == '{');
	++(*p);

	while (**p != '}') {
		len = read_key(p, buffer);
		assert(**p == ':');
		++(*p);

		if (memcmp("bufferView", buffer, len) == 0)
			accessor.view = read_uint(p);
		else if (memcmp("byteOffset", buffer, len) == 0)
			accessor.offset = read_uint(p);
		else if (memcmp("count", buffer, len) == 0)
			accessor.count = read_uint(p);
		else if (memcmp("sparse", buffer, len) == 0)
			assert(false); // Sparse accessors are not supported
		else
			skip_value(p);

		if (**p == ',')
			++(*p);
	}
	++(*p);

	return accessor;
}

static accessors_t read_accessors(const uint8_t *p, mesh_t mesh)
{
	accessors_t accessors = {};
	assert(*p == '[');
	++p;
	for (uint32_t i = 0; *p != ']'; ++i) {
		if (i == mesh.normals)
			accessors.normals = read_accessor(&p);
		else if (i == mesh.positions)
			accessors.positions = read_accessor(&p);
		else if (i == mesh.indices)
			accessors.indices = read_accessor(&p);
		else
			skip_value(&p);

		if (*p == ',')
			++p;
	}
	return accessors;
}

static view_t read_view(const uint8_t **p)
{
	char buffer[MAX_KEY_SIZE];
	size_t len;
	view_t view = {};

	assert(**p == '{');
	++(*p);

	while (**p != '}') {
		len = read_key(p, buffer);
		assert(**p == ':');
		++(*p);

		if (memcmp("buffer", buffer, len) == 0)
			assert(read_uint(p) == 0);
		else if (memcmp("byteOffset", buffer, len) == 0)
			view.offset = read_uint(p);
		else if (memcmp("byteLength", buffer, len) == 0)
			view.len = read_uint(p);
		else if (memcmp("byteStride", buffer, len) == 0)
			view.stride = read_uint(p);
		else
			skip_value(p);

		if (**p == ',')
			++(*p);
	}
	++(*p);

	return view;
}

static views_t read_views(const uint8_t *p, accessors_t accessors)
{
	views_t views = {};
	assert(*p == '[');
	++p;
	for (uint32_t i = 0; *p != ']'; ++i) {
		if (i == accessors.normals.view)
			views.normals = read_view(&p);
		else if (i == accessors.positions.view)
			views.positions = read_view(&p);
		else if (i == accessors.indices.view)
			views.indices = read_view(&p);
		else
			skip_value(&p);

		if (*p == ',')
			++p;
	}
	return views;
}

static gltf_info_t read_json(const uint8_t *json)
{
	char buffer[MAX_KEY_SIZE];
	size_t len;

	assert(*json == '{');
	++json;

	/*
	 * We want to avoid having to read every single accessor, buffer
	 * view etc, so we must identify which ones we need *before*
	 * processing these sections.  However, the sections may appear in
	 * any order, so we must first do a pass through the JSON to
	 * identify the section locations, so that we can then go through
	 * them in the necessary order.
	 */
	const uint8_t *meshes_json = nullptr;
	const uint8_t *accessors_json = nullptr;
	const uint8_t *views_json = nullptr;
	while (*json != '}') {
		len = read_key(&json, buffer);
		assert(*json == ':');
		++json;
		if (memcmp("meshes", buffer, len) == 0)
			meshes_json = json;
		else if (memcmp("accessors", buffer, len) == 0)
			accessors_json = json;
		else if (memcmp("bufferViews", buffer, len) == 0)
			views_json = json;
		skip_value(&json);
		if (*json == ',')
			++json;
	}
	assert(meshes_json != nullptr);
	assert(accessors_json != nullptr);
	assert(views_json != nullptr);

	mesh_t mesh = read_meshes(meshes_json);
	accessors_t accessors = read_accessors(accessors_json, mesh);
	views_t views = read_views(views_json, accessors);

	size_t normals_len, positions_len;
	if (accessors.positions.view == accessors.normals.view) {
		views.positions = views.normals;
		assert(views.normals.stride == sizeof(vec3_t));
		assert(views.positions.stride == sizeof(vec3_t));
		normals_len = accessors.normals.count * sizeof(vec3_t);
		positions_len = accessors.positions.count * sizeof(vec3_t);
	} else {
		normals_len = views.normals.len;
		positions_len = views.positions.len;
	}

	assert(views.indices.stride == 0);
	assert(views.indices.len / accessors.indices.count == sizeof(index_t));

	return (gltf_info_t) {
		.normals = {
			.offset = views.normals.offset + accessors.normals.offset,
			.len = normals_len,
		},
		.positions = {
			.offset = views.positions.offset + accessors.positions.offset,
			.len = positions_len,
		},
		.indices = {
			.offset = views.indices.offset + accessors.indices.offset,
			.len = views.indices.len,
		},
	};
}

static void load_vertices(
    vertex_t *restrict dest, const vec3_t *restrict positions,
    const vec3_t *restrict normals, size_t count)
{
	for (unsigned i = 0; i < count; ++i) {
		dest[i].position = positions[i];
		dest[i].normal = normals[i];
	}
}

model_t model_format_glb_read(const char *path)
{
	uint8_t *file = map_file(path);

	const file_hdr_t *file_hdr = (const file_hdr_t *)file;
	assert(file_hdr->magic == MAGIC);
	assert(file_hdr->version == VERSION);
	const size_t len = file_hdr->len;

	// Read JSON chunk
	size_t offset = sizeof(file_hdr_t);
	const chunk_hdr_t *json_hdr = (const chunk_hdr_t *)(file + offset);
	assert(json_hdr->type == CHUNK_TYPE_JSON);
	offset += sizeof(chunk_hdr_t);
	const gltf_info_t info = read_json(file + offset);

	// Check glTF info for validity
	assert(info.positions.len == info.normals.len);
	assert(info.positions.offset != info.normals.offset);
	if (info.positions.offset < info.normals.offset) {
		const uint32_t pos_end = info.positions.offset + info.positions.len;
		assert(pos_end <= info.normals.offset);
	} else {
		const uint32_t norm_end = info.normals.offset + info.normals.len;
		assert(norm_end <= info.positions.offset);
	}

	// Read binary chunk header
	offset += json_hdr->len;
	assert(offset != len);
	const chunk_hdr_t *bin_hdr = (const chunk_hdr_t *)(file + offset);
	assert(bin_hdr->type == CHUNK_TYPE_BIN);
	offset += sizeof(chunk_hdr_t);
	const uint8_t *bin_data = file + offset;
	offset += bin_hdr->len;
	assert(offset == len);

	// Load data from binary chunk into model
	const size_t vertex_count = info.positions.len / sizeof(vec3_t);
	const size_t index_count = info.indices.len / sizeof(index_t);
	model_t model = model_init(vertex_count, index_count);
	const vec3_t *positions
	    = (const vec3_t *)(bin_data + info.positions.offset);
	const vec3_t *normals = (const vec3_t *)(bin_data + info.normals.offset);
	load_vertices(model.vertices, positions, normals, vertex_count);
	memcpy(model.indices, bin_data + info.indices.offset, info.indices.len);

	munmap(file, len);
	return model;
}