#include "assets.h"
#include "file_formats/wavefront.h"
#include "platform.h"
#include <stdlib.h>
#include <string.h>
#include "stb_image.h"
#include "stb_vorbis.h"
#include <limits.h>
#include "debug/logger.h"
#include "render/render.h"
#include "cgltf.h"


Buffer read_file_into_buffer(const char *filename);



int wf_face_index_comp(const void *_a, const void *_b)
{
	wf_face_index *a = (wf_face_index *)_a;
	wf_face_index *b = (wf_face_index *)_b;

	if(a->vertex < b->vertex) return -1;
	if(a->vertex > b->vertex) return +1;
	if(a->texture < b->texture) return -1;
	if(a->texture > b->texture) return +1;
	if(a->normal < b->normal) return -1;
	if(a->normal > b->normal) return +1;
	return 0;
}

bool assets_load_object_obj(asset_manager *am, const char *filename, r_object *result)
{
	wf_obj_file obj_file;
	wf_mtl_file mtl_file;
	bool has_mtl;

	char path[512]; // @Buf: Files with paths longer than 512 bytes will fail to load
	p_file file;
	u64 size;
	Buffer content;
	bool success;

	// Load obj file
	strcpy(path, "assets/");
	strncat(path, filename, 512 - strlen(path));
	p_file_init(&file, path);
	size = p_file_size(&file);
	content.data = (u8*)p_alloc(size);
	content.size = 0;
	p_file_read(&file, &content, size);
	p_file_deinit(&file);

	success = wf_parse_obj((char*)content.data, content.size, filename, &obj_file);
	p_free(content.data);

	if(!success)
		return false;

	has_mtl = false;

	// Load mtl file (optional)
	if(obj_file.mtllib_names_count > 0)
	{
		has_mtl = true;
		char *mtl_name = obj_file.mtllib_names[0];

		strcpy(path, "assets/");
		strncat(path, mtl_name, 512 - strlen(path));
		p_file_init(&file, path);
		size = p_file_size(&file);
		content.data = (u8*)p_alloc(size);
		content.size = 0;
		p_file_read(&file, &content, size);
		p_file_deinit(&file);

		success = wf_parse_mtl((char*)content.data, content.size, mtl_name, &mtl_file);
		p_free(content.data);

		if(!success)
			return false;
	}



	// We take only the first object
	if(obj_file.objects_count == 0)
	{
		LOG(LOG_WARNING, "Asset: cannot load model. File does not contain any objects.");
		return false;
	}
	wf_object *object = obj_file.objects;

	// Setup model
	r_mesh *res_model = assets_new_models(am, 1);
	memset(res_model, 0, sizeof(r_mesh));
	// Actual conversion
	{
		wf_face_index *unique_indices;
		u32 unique_indices_count;

		unique_indices_count = object->faces_count * 3;
		unique_indices = (wf_face_index *)p_alloc(sizeof(wf_face_index) * unique_indices_count);
		memcpy(unique_indices, object->faces, sizeof(wf_face_index) * unique_indices_count);

		qsort(unique_indices, unique_indices_count, sizeof(wf_face_index), wf_face_index_comp);
		u32 removed = 0;
		for(u32 i = 1; i < unique_indices_count; i++)
		{
			wf_face_index *prev = unique_indices + i - 1;
			wf_face_index *curr = unique_indices + i;
			if(wf_face_index_comp(prev, curr) == 0) // Remove duplicate
				removed++;
			unique_indices[i - removed] = unique_indices[i];
		}
		unique_indices_count -= removed;

		res_model->vertices_count = unique_indices_count;
		res_model->vertices       = (v3*)p_alloc(sizeof(v3) * unique_indices_count);
		if(object->texture_coords)
			res_model->uvs = (v2*)p_alloc(sizeof(v2) * unique_indices_count);
		else
			res_model->uvs = NULL;
		res_model->normals        = (v3*)p_alloc(sizeof(v3) * unique_indices_count);
		for(u32 i = 0; i < unique_indices_count; i++)
		{
			wf_index vertex_i        = (unique_indices[i].vertex > 0  ? unique_indices[i].vertex - 1  : 0);
			wf_index texture_coord_i = (unique_indices[i].texture > 0 ? unique_indices[i].texture - 1 : 0);
			wf_index normal_i        = (unique_indices[i].normal > 0  ? unique_indices[i].normal - 1  : 0);

			res_model->vertices[i] = object->vertices [vertex_i].xyz / object->vertices[vertex_i].w;
			if(object->texture_coords)
				res_model->uvs[i] = object->texture_coords[texture_coord_i].uv;
			res_model->normals[i] = object->normals[normal_i];
		}

		res_model->indices_count = object->faces_count * 3;
		res_model->indices = (u32*)p_alloc(sizeof(u32) * object->faces_count * 3);
		for(u32 f = 0; f < object->faces_count; f++)
		{
			wf_face *curr_face = object->faces + f;
			for(u32 fi = 0; fi < 3; fi++)
			{
				wf_face_index *curr_index = curr_face->indices + fi;
				wf_face_index *found_ptr = (wf_face_index *)bsearch(curr_index, unique_indices, unique_indices_count, sizeof(wf_face_index), wf_face_index_comp);
				res_model->indices[3 * f + fi] = (found_ptr - unique_indices);
			}
		}

		p_free(unique_indices);
	}

	// Setup material
	r_material *res_material = assets_new_materials(am, 1);
	memset(res_material, 0, sizeof(r_material));
	res_material->shader = &r_render_state.shader_pbr;

	// Load textures
	if(object->material_name != NULL && has_mtl)
	{
		// @Feature: Load multiple materials
		wf_material *material = NULL;
		for(wf_size i = 0; i < mtl_file.materials_count; i++)
		{
			if(strcmp(object->material_name, mtl_file.materials[i].name) == 0)
			{
				material = mtl_file.materials + i;
				break;
			}
		}

		if(material)
		{
			// Load diffuse texture
			if(material->map_Kd && strcmp(material->map_Kd, "") != 0)
			{
				strcpy(path, "assets/");
				strncat(path, material->map_Kd, 512 - strlen(path));

				stbi_set_flip_vertically_on_load(true);
				v2s size; s32 channels;
				u8 *data = stbi_load(path, &size.x, &size.y, &channels, 4);

				r_texture *texture = assets_new_textures(am, 1);
				*texture = r_texture_create(data, size, R_TEXTURE_SRGB);
				res_material->albedo_texture = texture;
				res_material->albedo_factor  = v4{1,1,1,1};
			}

			// Load emissive texture
			if(material->map_Ke && strcmp(material->map_Ke, "") != 0)
			{
				strcpy(path, "assets/");
				strncat(path, material->map_Ke, 512 - strlen(path));

				stbi_set_flip_vertically_on_load(true);
				v2s size; s32 channels;
				u8 *data = stbi_load(path, &size.x, &size.y, &channels, 4);

				r_texture *texture = assets_new_textures(am, 1);
				*texture = r_texture_create(data, size, R_TEXTURE_SRGB);
				res_material->emissive_texture = texture;
				res_material->emissive_factor  = v4{1,1,1,1};
			}
		}
	}

	result->meshes = (r_mesh**)p_alloc(sizeof(void*));
	result->mesh_material = (r_material**)p_alloc(sizeof(void*));
	result->meshes[0] = res_model;
	result->mesh_material[0] = res_material;
	result->mesh_local_transform = (m4*)p_alloc(sizeof(m4));
	result->mesh_local_transform[0] = m4_identity;
	result->count = 1;
	result->scale = v3{1,1,1};
	result->position = v3{0,0,0};
	result->rotation = v3{0,0,0};

	wf_cleanup_obj(&obj_file);
	if(has_mtl)
	{
		wf_cleanup_mtl(&mtl_file);
	}
	return true;
}


struct gltf_loader_state
{
	asset_manager *am;
	cgltf_data *data;

	r_texture *texture_buffer;
	u32 texture_capacity;
};

void gltf_loader_reserve_texture_space(gltf_loader_state *state, u32 count)
{
	state->texture_capacity = count;
	state->texture_buffer = assets_new_textures(state->am, state->texture_capacity);
	memset(state->texture_buffer, 0, sizeof(r_texture) * state->texture_capacity);
}

r_texture *gltf_lazy_load_texture(gltf_loader_state *state, u32 texture_index, bool sRGB)
{
	assert(texture_index < state->texture_capacity);

	if(state->texture_buffer[texture_index].flags & R_TEXTURE_INITIALIZED)
	{
		bool is_buffered_sRGB = !!(state->texture_buffer[texture_index].flags & (sRGB ? R_TEXTURE_SRGB : 0));
		assert(sRGB == is_buffered_sRGB);
		return &state->texture_buffer[texture_index];
	}

	cgltf_data *data = state->data;
	cgltf_buffer_view *t_view = data->images[texture_index].buffer_view;

	LOG(LOG_DEBUG, "Loading texture %s, sRGB: %d", (t_view->name ? t_view->name : t_view->buffer->name), sRGB);
	u8 *t_data;
	s32 t_width, t_height, t_channels;
	stbi_set_flip_vertically_on_load(false);
	t_data = stbi_load_from_memory((u8*)t_view->buffer->data + t_view->offset, t_view->size, &t_width, &t_height, &t_channels, 4);

	u32 flags = (sRGB ? R_TEXTURE_SRGB : 0);
	state->texture_buffer[texture_index] = r_texture_create(t_data, v2s{t_width, t_height}, flags);
	return &state->texture_buffer[texture_index];
}

void gltf_visit(gltf_loader_state *state, r_object *object, m4 transform, cgltf_node *node)
{
	cgltf_data *data = state->data;

	// Transform
	if(node->has_matrix)
	{
		transform = transpose(*(m4*)node->matrix) * transform;
	}
	else
	{
		if(node->has_translation)
		{
			v3 t = {node->translation[0], node->translation[1], node->translation[2]};
			transform = translation_v3(t) * transform;
		}
		if(node->has_rotation)
		{
			LOG(LOG_ERROR, "Quaternion rotation not supported yet");
		}
		if(node->has_scale)
		{
			v3 s = {node->scale[0], node->scale[1], node->scale[2]};
			transform = scale_v3(s) * transform;
		}
	}

	// @Feature: Skin
	if(node->mesh)
	{
		cgltf_mesh *mesh = node->mesh;
		for(cgltf_size p = 0; p < mesh->primitives_count; p++)
		{
			cgltf_primitive *primitive = &mesh->primitives[p];

			if(primitive->type == cgltf_primitive_type_triangles)
			{
				// Mesh
				v3  *vertices = NULL;
				v3  *normals  = NULL;
				v3  *tangents = NULL;
				v2  *uvs      = NULL;
				u64 vertices_count = 0;

				u32 *indices = NULL;
				u64 indices_count = 0;

				if(primitive->indices->type != cgltf_type_scalar)
				{
					(LOG_ERROR, "glTF: Unhandled indices type %d", primitive->indices->type);
					continue;
				}

				cgltf_accessor *a = primitive->indices;
				indices = (u32*)p_alloc(a->count * sizeof(u32));
				indices_count = a->count;
				cgltf_accessor_unpack_indices(a, indices, a->count);

				vertices_count = primitive->attributes[0].data->count;
				for(cgltf_size i = 0; i < primitive->attributes_count; i++)
				{
					if(primitive->attributes[i].type == cgltf_attribute_type_position)
					{
						cgltf_accessor *a = primitive->attributes[i].data;
						vertices = (v3*)p_alloc(a->count * sizeof(v3));
						if(vertices_count != a->count)
						{
							LOG(LOG_ERROR, "glTF: vertex count (%d) != previously stored vertex count (%d)", a->count, vertices_count);
							continue;
						}
						cgltf_accessor_unpack_floats(a, (f32*)vertices, a->count * 3);
					}
					else if(primitive->attributes[i].type == cgltf_attribute_type_normal)
					{
						cgltf_accessor *a = primitive->attributes[i].data;
						normals = (v3*)p_alloc(a->count * sizeof(v3));
						if(vertices_count != a->count)
						{
							LOG(LOG_ERROR, "glTF: normal count (%d) != vertex count (%d)", a->count, vertices_count);
							continue;
						}
						cgltf_accessor_unpack_floats(a, (f32*)normals, a->count * 3);
					}
					else if(primitive->attributes[i].type == cgltf_attribute_type_tangent)
					{
						cgltf_accessor *a = primitive->attributes[i].data;
						tangents         = (v3*)p_alloc(a->count * sizeof(v3));
						v4 *tangents_tmp = (v4*)p_alloc(a->count * sizeof(v4));
						if(vertices_count != a->count)
						{
							LOG(LOG_ERROR, "glTF: tangent count (%d) != vertex count (%d)", a->count, vertices_count);
							continue;
						}
						cgltf_accessor_unpack_floats(a, (f32*)tangents_tmp, a->count * 4);
						for(u32 i = 0; i < a->count; i++)
							tangents[i] = tangents_tmp[i].xyz;
						p_free(tangents_tmp);
					}
					else if(primitive->attributes[i].type == cgltf_attribute_type_texcoord)
					{
						cgltf_accessor *a = primitive->attributes[i].data;
						uvs = (v2*)p_alloc(a->count * sizeof(v2));
						if(vertices_count != a->count)
						{
							LOG(LOG_ERROR, "glTF: uv count (%d) != vertex count (%d)", a->count, vertices_count);
							continue;
						}
						cgltf_accessor_unpack_floats(a, (f32*)uvs, a->count * 2);
					}
					else
					{
						//LOG(LOG_WARNING, "Unmanaged attribute type: %d", primitive->attributes[i].type);
					}
				}

				r_mesh *mesh     = (r_mesh*)p_alloc(sizeof(r_mesh));
				*mesh = r_mesh_create(indices_count, indices, vertices_count, vertices, normals, tangents, uvs);

				// @Feature: Calculate tangents if missing


				// Material
				r_material *material = assets_new_materials(state->am, 1);
				memset(material, 0, sizeof(r_material));
				material->shader = &r_render_state.shader_pbr;

				// Find image index
				u32 albedo_texture_index = primitive->material->pbr_metallic_roughness.base_color_texture.texture->image - data->images;
				if(albedo_texture_index > data->images_count)
					LOG(LOG_ERROR, "Albedo image index (%u) out of bounds (%u)", albedo_texture_index, data->images_count);

				material->albedo_texture = gltf_lazy_load_texture(state, albedo_texture_index, true);
				material->albedo_factor  = V4(primitive->material->pbr_metallic_roughness.base_color_factor);


				// @Cleanup: @Performance: Merge metallic and roughness texture or split them up?
				u32 metallic_texture_index = primitive->material->pbr_metallic_roughness.metallic_roughness_texture.texture->image - data->images;
				if(metallic_texture_index > data->images_count)
					LOG(LOG_ERROR, "Metallic image index (%u) out of bounds (%u)", metallic_texture_index, data->images_count);

				material->metallic_texture = gltf_lazy_load_texture(state, metallic_texture_index, false);
				material->metallic_factor  = primitive->material->pbr_metallic_roughness.metallic_factor;



				u32 roughness_texture_index = primitive->material->pbr_metallic_roughness.metallic_roughness_texture.texture->image - data->images;
				if(roughness_texture_index > data->images_count)
					LOG(LOG_ERROR, "Roughness image index (%u) out of bounds (%u)", roughness_texture_index, data->images_count);

				material->roughness_texture = gltf_lazy_load_texture(state, roughness_texture_index, false);
				material->roughness_factor  = primitive->material->pbr_metallic_roughness.roughness_factor;


				if(primitive->material->normal_texture.texture)
				{
					u32 normal_texture_index = primitive->material->normal_texture.texture->image - data->images;
					if(normal_texture_index > data->images_count)
						LOG(LOG_ERROR, "Normal image index (%u) out of bounds (%u)", normal_texture_index, data->images_count);

					material->normal_texture = gltf_lazy_load_texture(state, normal_texture_index, false);
				}


				if(primitive->material->emissive_texture.texture)
				{
					u32 emissive_texture_index = primitive->material->emissive_texture.texture->image - data->images;
					if(emissive_texture_index > data->images_count)
						LOG(LOG_ERROR, "Emission image index (%u) out of bounds (%u)", emissive_texture_index, data->images_count);

					material->emissive_texture = gltf_lazy_load_texture(state, emissive_texture_index, true);
					f32 strength = primitive->material->has_emissive_strength ? primitive->material->has_emissive_strength : 1.0;
					material->emissive_factor  = V4(V3(primitive->material->emissive_factor) * strength, 1.0);
				}


				object->meshes        = (r_mesh**)p_realloc(object->meshes, (object->count + 1) * sizeof(r_mesh));
				object->mesh_material = (r_material**)p_realloc(object->mesh_material, (object->count + 1) * sizeof(r_material));
				object->mesh_local_transform = (m4*)p_realloc(object->mesh_local_transform, (object->count + 1) * sizeof(m4));

				object->meshes   [object->count] = mesh;
				object->mesh_material[object->count] = material;
				object->mesh_local_transform[object->count] = transform;
				object->count++;
			}
			else
			{
				LOG(LOG_ERROR, "glTF: Unhandled primitive type %d", primitive->type);
			}
		}
	}


	for(cgltf_size i = 0; i < node->children_count; i++)
	{
		gltf_visit(state, object, transform, node->children[i]);
	}
}

bool assets_load_object_gltf(asset_manager *am, const char *filename, r_object *result)
{
	gltf_loader_state state;
	state.am = am;

	cgltf_options options = {};
	cgltf_data *data = NULL;

	cgltf_result status = cgltf_parse_file(&options, filename, &data);
	if(status != cgltf_result_success)
	{
		LOG(LOG_DEBUG, "Error parsing glTF file \"%s\".", filename);
		return false;
	}

	status = cgltf_load_buffers(&options, data, filename);
	if(status != cgltf_result_success)
	{
		LOG(LOG_DEBUG, "Erorr loading glTF buffers from file \"%s\".", filename);
		cgltf_free(data);
		return false;
	}
	state.data = data;

	// Reserve texture space
	gltf_loader_reserve_texture_space(&state, data->images_count);

	// Prepare object
	r_object object;
	memset(&object, 0, sizeof(r_object));
	object.scale = v3{1,1,1};
	object.position = v3{0,0,0};
	object.rotation = v3{0,0,0};
	object.has_shadow = true;

	for(cgltf_size i = 0; i < data->scene->nodes_count; i++)
	{
		gltf_visit(&state, &object, m4_identity, data->scene->nodes[i]);
	}

	// Resize unused space

	*result = object;
	cgltf_free(data);
	return true;
}

bool assets_load_audio(asset_manager *am, const char *name, p_audio_buffer *result)
{
	s32 music_channels;
	s32 music_sample_rate;
	s16 *music_data;
	s32 music_samples = stb_vorbis_decode_filename(name, &music_channels, &music_sample_rate, &music_data);

	result->samples = (p_audio_sample*) p_alloc(sizeof(p_audio_sample) * music_samples);
	for(u32 i = 0; i < music_samples; i++)
	{
		result->samples[i].left  = (f32)music_data[2*i]     / -SHRT_MIN; // Left
		result->samples[i].right = (f32)music_data[2*i + 1] / -SHRT_MIN; // Right
	}
	result->size = music_samples;

	free(music_data);

	return true;
}


bool assets_load_cubemap(asset_manager *am, const char *px, const char *nx, const char *py, const char *ny, const char *pz, const char *nz, r_cubemap *result)
{
	stbi_set_flip_vertically_on_load(false);
	//stbi_ldr_to_hdr_scale(1.0f);
	//stbi_ldr_to_hdr_gamma(2.2f);

	float *data[6];
	v2s size; s32 channels;

	v2s s;
	data[0] = stbi_loadf(px, &s.x, &s.y, &channels, 3);
	size = s;
	data[1] = stbi_loadf(nx, &s.x, &s.y, &channels, 3);
	assert(size == s);
	data[2] = stbi_loadf(py, &s.x, &s.y, &channels, 3);
	assert(size == s);
	data[3] = stbi_loadf(ny, &s.x, &s.y, &channels, 3);
	assert(size == s);
	data[4] = stbi_loadf(pz, &s.x, &s.y, &channels, 3);
	assert(size == s);
	data[5] = stbi_loadf(nz, &s.x, &s.y, &channels, 3);
	assert(size == s);

	r_cubemap cubemap = r_cubemap_create(data, size, 0);
	if(result)
		*result = cubemap;

	return true;
}


void assets_init(asset_manager *am)
{
	am->shaders = NULL;
	am->shaders_count = 0;

	am->models = NULL;
	am->models_count = 0;

	am->textures = NULL;
	am->textures_count = 0;

	am->materials = NULL;
	am->materials_count = 0;

	am->objects = NULL;
	am->objects_count = 0;

	am->allocations = NULL;
	am->allocations_count = 0;
}

void assets_free(asset_manager *am)
{

}

static void assets_add_allocation(asset_manager *am, void *data)
{
	u32 old_count = am->allocations_count;
	am->allocations_count++;
	am->allocations = (void**)p_realloc(am->allocations, am->allocations_count * sizeof(void*));
	am->allocations[old_count] = data;
}


r_shader *assets_new_shaders(asset_manager *am, u32 count)
{
	u32 old_count = am->shaders_count;
	am->shaders_count = old_count + count;

	r_shader *data = (r_shader*) p_alloc(count * sizeof(r_shader));
	assets_add_allocation(am, data);

	am->shaders = (r_shader**) p_realloc(am->shaders, am->shaders_count * sizeof(r_shader*));
	for(u32 i = 0; i < count; i++)
		am->shaders[old_count + i] = data + i;

	return data;
}

r_mesh *assets_new_models(asset_manager *am, u32 count)
{
	u32 old_count = am->models_count;
	am->models_count = old_count + count;

	r_mesh *data = (r_mesh*) p_alloc(count * sizeof(r_mesh));
	assets_add_allocation(am, data);

	am->models = (r_mesh**) p_realloc(am->models, am->models_count * sizeof(r_mesh*));
	for(u32 i = 0; i < count; i++)
		am->models[old_count + i] = data + i;

	return data;
}

r_texture *assets_new_textures(asset_manager *am, u32 count)
{
	u32 old_count = am->textures_count;
	am->textures_count = old_count + count;

	r_texture *data = (r_texture*) p_alloc(count * sizeof(r_texture));
	assets_add_allocation(am, data);

	am->textures = (r_texture**) p_realloc(am->textures, am->textures_count * sizeof(r_texture*));
	for(u32 i = 0; i < count; i++)
		am->textures[old_count + i] = data + i;

	return data;
}

r_material *assets_new_materials(asset_manager *am, u32 count)
{
	u32 old_count = am->materials_count;
	am->materials_count = old_count + count;

	r_material *data = (r_material*) p_alloc(count * sizeof(r_material));
	assets_add_allocation(am, data);

	am->materials = (r_material**) p_realloc(am->materials, am->materials_count * sizeof(r_material*));
	for(u32 i = 0; i < count; i++)
		am->materials[old_count + i] = data + i;

	return data;
}

r_object *assets_new_objects(asset_manager *am, u32 count)
{
	u32 old_count = am->objects_count;
	am->objects_count = old_count + count;

	r_object *data = (r_object*) p_alloc(count * sizeof(r_object));
	assets_add_allocation(am, data);

	am->objects = (r_object**) p_realloc(am->objects, am->objects_count * sizeof(r_object*));
	for(u32 i = 0; i < count; i++)
		am->objects[old_count + i] = data + i;

	return data;
}










Buffer read_file_into_buffer(const char *filename)
{
	Buffer empty = { .size = 0, .data = NULL };
	Buffer result;
	p_file file;
	bool success;

	success = p_file_init(&file, filename);
	if(success) {
		result.size = p_file_size(&file);
		result.data = (u8*)p_alloc(result.size);
		success = p_file_read(&file, &result, result.size);

		p_file_deinit(&file);
	}

	if(!success)
		return empty;

	return result;
}