#include "model_format_glb.h"

#include <assert.h>
#include <ctype.h>
#include <fcntl.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.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_STRING_SIZE 128
#define MAX_JOINT_COUNT 32

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

enum {
	RENDERING_MODE_TRIANGLES = 4,
};

typedef enum {
	COMPONENT_TYPE_S8 = 5120, // Byte
	COMPONENT_TYPE_U8 = 5121, // Unsigned byte
	COMPONENT_TYPE_S16 = 5122, // Short
	COMPONENT_TYPE_U16 = 5123, // Unsigned short
	COMPONENT_TYPE_U32 = 5125, // Unsigned int
	COMPONENT_TYPE_F32 = 5126, // Float
} component_type_t;

typedef enum {
	ACCESSOR_TYPE_SCALAR,
	ACCESSOR_TYPE_VEC2,
	ACCESSOR_TYPE_VEC3,
	ACCESSOR_TYPE_VEC4,
	ACCESSOR_TYPE_MAT2,
	ACCESSOR_TYPE_MAT3,
	ACCESSOR_TYPE_MAT4,
} accessor_type_t;

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 bone_indices;
	uint32_t weights;
	uint32_t indices;
} mesh_t;

typedef struct {
	uint32_t inverse_bind_matrices;
	uint32_t joints[MAX_JOINT_COUNT];
	uint32_t joint_count;
} skin_t;

typedef struct {
	uint32_t view;
	uint32_t offset;
	uint32_t count;
	component_type_t component;
	accessor_type_t type;
} accessor_t;

typedef struct {
	accessor_t normals;
	accessor_t positions;
	accessor_t bone_indices;
	accessor_t weights;
	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 bone_indices;
	view_t weights;
	view_t indices;
} views_t;

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

typedef struct {
	bone_t *bones;
	gltf_info_member_t normals;
	gltf_info_member_t positions;
	gltf_info_member_t bone_indices;
	gltf_info_member_t weights;
	gltf_info_member_t indices;
	size_t vertex_count, index_count, bone_count;
} 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 void read_string(const uint8_t **p, char buffer[MAX_STRING_SIZE])
{
	assert(**p == '"');
	++(*p);
	size_t len = 0;
	while (**p != '"') {
		assert(len < MAX_STRING_SIZE - 1);
		buffer[len++] = *(*p)++;
	}
	++(*p);
	buffer[len] = '\0';
}

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 float read_float(const uint8_t **p)
{
	enum {
		INTEGRAL,
		DECIMAL,
		EXPONENT,
	} state
	    = INTEGRAL;

	bool neg_val = false, neg_exp = false;
	float val = 0, exp = 0;
	unsigned div = 10;
	while (!ends_value(**p)) {
		const uint8_t c = *((*p)++);
		switch (state) {
		case INTEGRAL:
			if (isdigit(c))
				val = 10 * val + (c - '0');
			else if (c == '-')
				neg_val = true;
			else if (c == '.')
				state = DECIMAL;
			else
				assert(false);
			break;

		case DECIMAL:
			if (isdigit(c))
				val += (float)(c - '0') / div;
			else if (c == 'e')
				state = EXPONENT;
			else
				assert(false);
			div *= 10;
			break;

		case EXPONENT:
			if (isdigit(c))
				exp = 10 * exp + (c - '0');
			else if (c == '-')
				neg_exp = true;
			else
				assert(false);
			break;

		default:
			assert(false);
		}
	}

	if (neg_val)
		val *= -1;
	if (neg_exp)
		exp *= -1;
	return val * powf(10, exp);
}

static vec3_t read_vec3(const uint8_t **p)
{
	vec3_t v;
	assert(**p == '[');
	++(*p);
	v.x = read_float(p);
	assert(**p == ',');
	++(*p);
	v.y = read_float(p);
	assert(**p == ',');
	++(*p);
	v.z = read_float(p);
	assert(**p == ']');
	++(*p);
	return v;
}

static vec4_t read_vec4(const uint8_t **p)
{
	vec4_t v;
	assert(**p == '[');
	++(*p);
	v.x = read_float(p);
	assert(**p == ',');
	++(*p);
	v.y = read_float(p);
	assert(**p == ',');
	++(*p);
	v.z = read_float(p);
	assert(**p == ',');
	++(*p);
	v.w = read_float(p);
	assert(**p == ']');
	++(*p);
	return v;
}

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

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

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

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

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

			while (**p != '}') {
				read_string(p, buffer);
				assert(**p == ':');
				++(*p);
				if (strcmp("NORMAL", buffer) == 0) {
					mesh.normals = read_uint(p);
				} else if (strcmp("POSITION", buffer) == 0) {
					mesh.positions = read_uint(p);
				} else if (strcmp("JOINTS_0", buffer) == 0) {
					mesh.bone_indices = read_uint(p);
				} else if (strcmp("WEIGHTS_0", buffer) == 0) {
					mesh.weights = read_uint(p);
				} else {
					skip_value(p);
				}
				if (**p == ',')
					++(*p);
			}
			++(*p);
		} else if (strcmp("indices", buffer) == 0) {
			mesh.indices = read_uint(p);
		} else if (strcmp("mode", buffer) == 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_STRING_SIZE];

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

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

	// Primitives key missing from mesh object
	assert(false);
}

static skin_t read_skins(const uint8_t *p)
{
	skin_t skin = {};
	char buffer[MAX_STRING_SIZE];

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

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

		if (strcmp("inverseBindMatrices", buffer) == 0) {
			skin.inverse_bind_matrices = read_uint(&p);
		} else if (strcmp("joints", buffer) == 0) {
			assert(*p == '[');
			++p;
			for (skin.joint_count = 0; *p != ']'; ++skin.joint_count) {
				assert(skin.joint_count < MAX_JOINT_COUNT);
				skin.joints[skin.joint_count] = read_uint(&p);

				if (*p == ',')
					++p;
			}
			++p;
		} else {
			skip_value(&p);
		}

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

	assert(*p == ']');
	return skin;
}

static void read_joint_children(
    const uint8_t **p, skin_t skin, bone_t bones[MAX_JOINT_COUNT],
    uint32_t parent)
{
	assert(**p == '[');
	++(*p);

	while (**p != ']') {
		const uint32_t i = read_uint(p);
		uint32_t j;
		for (j = 0; j < skin.joint_count; ++j) {
			if (skin.joints[j] == i)
				break;
		}
		assert(j < skin.joint_count);
		assert(bones[j].parent == UINT32_MAX);
		bones[j].parent = parent;

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

static void
read_joints(const uint8_t *p, skin_t skin, bone_t bones[MAX_JOINT_COUNT])
{
	char buffer[MAX_STRING_SIZE];

	for (uint32_t i = 0; i < skin.joint_count; ++i)
		bones[i].parent = UINT32_MAX;

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

	for (uint32_t i = 0; *p != ']'; ++i) {
		uint32_t j;
		for (j = 0; j < skin.joint_count; ++j) {
			if (skin.joints[j] == i)
				break;
		}
		if (j >= skin.joint_count) {
			skip_value(&p);
			if (*p == ',')
				++p;
			continue;
		}

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

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

			if (strcmp("children", buffer) == 0)
				read_joint_children(&p, skin, bones, j);
			else if (strcmp("translation", buffer) == 0)
				bones[j].translation = read_vec3(&p);
			else if (strcmp("rotation", buffer) == 0)
				bones[j].rotation = read_vec4(&p);
			else if (strcmp("scale", buffer) == 0)
				bones[j].scale = read_vec3(&p);
			else
				skip_value(&p);

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

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

static component_type_t read_component_type(const uint8_t **p)
{
	const uint32_t component = read_uint(p);
	assert(component >= COMPONENT_TYPE_S8);
	assert(component <= COMPONENT_TYPE_F32);
	return (component_type_t)component;
}

static accessor_type_t read_accessor_type(const uint8_t **p)
{
	char buffer[MAX_STRING_SIZE];
	read_string(p, buffer);
	if (strcmp("SCALAR", buffer) == 0)
		return ACCESSOR_TYPE_SCALAR;
	else if (strcmp("VEC3", buffer) == 0)
		return ACCESSOR_TYPE_VEC3;
	else if (strcmp("VEC4", buffer) == 0)
		return ACCESSOR_TYPE_VEC4;
	else if (strcmp("MAT4", buffer) == 0)
		return ACCESSOR_TYPE_MAT4;
	else
		assert(false);
}

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

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

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

		if (strcmp("bufferView", buffer) == 0)
			accessor.view = read_uint(p);
		else if (strcmp("byteOffset", buffer) == 0)
			accessor.offset = read_uint(p);
		else if (strcmp("count", buffer) == 0)
			accessor.count = read_uint(p);
		else if (strcmp("componentType", buffer) == 0)
			accessor.component = read_component_type(p);
		else if (strcmp("type", buffer) == 0)
			accessor.type = read_accessor_type(p);
		else if (strcmp("sparse", buffer) == 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.bone_indices)
			accessors.bone_indices = read_accessor(&p);
		else if (i == mesh.weights)
			accessors.weights = 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_STRING_SIZE];
	view_t view = {};

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

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

		if (strcmp("buffer", buffer) == 0)
			assert(read_uint(p) == 0);
		else if (strcmp("byteOffset", buffer) == 0)
			view.offset = read_uint(p);
		else if (strcmp("byteLength", buffer) == 0)
			view.len = read_uint(p);
		else if (strcmp("byteStride", buffer) == 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.bone_indices.view)
			views.bone_indices = read_view(&p);
		else if (i == accessors.weights.view)
			views.weights = read_view(&p);
		else if (i == accessors.indices.view)
			views.indices = read_view(&p);
		else
			skip_value(&p);

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

	if (accessors.normals.view == accessors.positions.view)
		views.positions = views.normals;

	assert(views.normals.len != 0);
	assert(views.positions.len != 0);
	assert(views.bone_indices.len != 0);
	assert(views.weights.len != 0);
	assert(views.indices.len != 0);
	return views;
}

static gltf_info_member_t info_member(accessor_t accessor, view_t view)
{
	return (gltf_info_member_t) {
		.len = view.stride == 0 ? view.len : accessor.count * view.stride,
		.offset = view.offset + accessor.offset,
	};
}

static gltf_info_t read_json(const uint8_t *json)
{
	char buffer[MAX_STRING_SIZE];

	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 *skins_json = nullptr;
	const uint8_t *nodes_json = nullptr;
	const uint8_t *accessors_json = nullptr;
	const uint8_t *views_json = nullptr;
	while (*json != '}') {
		read_string(&json, buffer);
		assert(*json == ':');
		++json;
		if (strcmp("meshes", buffer) == 0)
			meshes_json = json;
		else if (strcmp("skins", buffer) == 0)
			skins_json = json;
		else if (strcmp("nodes", buffer) == 0)
			nodes_json = json;
		else if (strcmp("accessors", buffer) == 0)
			accessors_json = json;
		else if (strcmp("bufferViews", buffer) == 0)
			views_json = json;
		skip_value(&json);
		if (*json == ',')
			++json;
	}
	assert(meshes_json != nullptr);
	assert(skins_json != nullptr);
	assert(nodes_json != nullptr);
	assert(accessors_json != nullptr);
	assert(views_json != nullptr);

	const mesh_t mesh = read_meshes(meshes_json);
	const skin_t skin = read_skins(skins_json);
	bone_t *bones = calloc(skin.joint_count, sizeof(bone_t));
	assert(bones != nullptr);
	read_joints(nodes_json, skin, bones);

	const accessors_t accessors = read_accessors(accessors_json, mesh);
	const views_t views = read_views(views_json, accessors);

	// Verify data types
	assert(
	    accessors.normals.component == COMPONENT_TYPE_F32
	    && accessors.normals.type == ACCESSOR_TYPE_VEC3);
	assert(
	    accessors.positions.component == COMPONENT_TYPE_F32
	    && accessors.positions.type == ACCESSOR_TYPE_VEC3);
	assert(
	    accessors.bone_indices.component == COMPONENT_TYPE_U16
	    && accessors.bone_indices.type == ACCESSOR_TYPE_VEC4);
	assert(
	    accessors.weights.component == COMPONENT_TYPE_F32
	    && accessors.weights.type == ACCESSOR_TYPE_VEC4);
	assert(
	    accessors.indices.component == COMPONENT_TYPE_U16
	    && accessors.indices.type == ACCESSOR_TYPE_SCALAR);

	// Ensure the vertex count is consistent
	assert(
	    accessors.positions.count == accessors.normals.count
	    && accessors.positions.count == accessors.bone_indices.count
	    && accessors.positions.count == accessors.weights.count);

	return (gltf_info_t) {
		.bones = bones,
		.positions = info_member(accessors.positions, views.positions),
		.normals = info_member(accessors.normals, views.normals),
		.bone_indices
		= info_member(accessors.bone_indices, views.bone_indices),
		.weights = info_member(accessors.weights, views.weights),
		.indices = info_member(accessors.indices, views.indices),
		.vertex_count = accessors.positions.count,
		.index_count = accessors.indices.count,
		.bone_count = skin.joint_count,
	};
}

static void load_vertices(
    vertex_t *restrict dest, const vec3_t *restrict positions,
    const vec3_t *restrict normals, const u16_vec4_t *restrict bones,
    const vec4_t *restrict weights, size_t count)
{
	for (unsigned i = 0; i < count; ++i) {
		dest[i].position = positions[i];
		dest[i].normal = normals[i];
		dest[i].bones = bones[i];
		dest[i].weights = weights[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);

	// 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
	model_t model = model_init(info.vertex_count, info.index_count);
	model.bone_count = info.bone_count;
	model.bones = info.bones;
	load_vertices(
	    model.vertices, (const vec3_t *)(bin_data + info.positions.offset),
	    (const vec3_t *)(bin_data + info.normals.offset),
	    (const u16_vec4_t *)(bin_data + info.bone_indices.offset),
	    (const vec4_t *)(bin_data + info.weights.offset), info.vertex_count);
	memcpy(model.indices, bin_data + info.indices.offset, info.indices.len);

	munmap(file, len);
	return model;
}