#include "model_format_glb.h"

#include <assert.h>
#include <ctype.h>
#include <fcntl.h>
#include <math.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_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 bones;
	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 parent;
	vec3_t translation;
	vec4_t rotation;
	vec3_t scale;
} joint_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 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 {
					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, joint_t joints[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(joints[j].parent == UINT32_MAX);
		joints[j].parent = parent;

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

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

	for (uint32_t i = 0; i < skin.joint_count; ++i)
		joints[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, joints, j);
			else if (strcmp("translation", buffer) == 0)
				joints[j].translation = read_vec3(&p);
			else if (strcmp("rotation", buffer) == 0)
				joints[j].rotation = read_vec4(&p);
			else if (strcmp("scale", buffer) == 0)
				joints[j].scale = read_vec3(&p);
			else
				skip_value(&p);

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

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

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, component_type_t component_type, accessor_type_t type)
{
	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)
			assert(read_uint(p) == component_type);
		else if (strcmp("type", buffer) == 0)
			assert(read_accessor_type(p) == type);
		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, COMPONENT_TYPE_F32, ACCESSOR_TYPE_VEC3);
		} else if (i == mesh.positions) {
			accessors.positions
			    = read_accessor(&p, COMPONENT_TYPE_F32, ACCESSOR_TYPE_VEC3);
		} else if (i == mesh.indices) {
			accessors.indices = read_accessor(
			    &p, COMPONENT_TYPE_U16, ACCESSOR_TYPE_SCALAR);
		} 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.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;
	return views;
}

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);
	joint_t joints[MAX_JOINT_COUNT] = {};
	read_joints(nodes_json, skin, joints);

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

	// Assume indices are in their own buffer view
	assert(views.indices.stride == 0);

	return (gltf_info_t) {
		.normals = {
			.offset = views.normals.offset + accessors.normals.offset,
			.len = accessors.normals.count * sizeof(vec3_t),
		},
		.positions = {
			.offset = views.positions.offset + accessors.positions.offset,
			.len = accessors.positions.count * sizeof(vec3_t),
		},
		.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;
}