module dlib.assets;

import includes;
import dlib.aliases;
import dlib.util;
import dlib.alloc;
import dlib.math;

import std.file;
import std.path;
import std.format;
import std.stdio;
import std.exception;

__gshared ImageData DEFAULT_IMAGE    = GenerateDefaultTexture(32, 32);
__gshared Material  DEFAULT_MATERIAL = GenerateDefaultMaterial();

// TODO:
// Alignment?

/*************************
 ****** FILE PACKER ******
 *************************/

/*************************
 ****** FILE PACKER ******
 *************************/

/************************
 *** ASSET STRUCTURES ***
 ************************/

struct ImageData
{
	u8[] data;
	u32  w;
	u32  h;
	u32  ch;
}

struct ModelData
{
	Vertex[]    vtx_buf;
	u32[]       idx_buf;
	Mesh[]      meshes;
	Material[]  mats;
	ImageData[] tex;
}

/************************
 *** ASSET STRUCTURES ***
 ************************/

/**************************
 ***** GFX STRUCTURES *****
 **************************/

enum MaterialMapIndex
{
	Albedo,
	Normal,
	Metallic,
	Roughness,
	Occlusion,
	Emission,
	Height,
	Cubemap,
	Irradiance,
	Prefilter,
	BRDF,

	Max,
}

alias MMI = MaterialMapIndex;

struct MaterialMap
{
	Vec4 col;
	u32  tex_id;
	f32  value;
}

struct Material
{
	u32                  pipeline_id;
	MaterialMap[MMI.max] maps;
}

struct Vertex
{
	Vec4    color;
	Vec4    tangent;
	Vec3    pos;
	Vec3    normal;
	Vec2[2] uv;
}

struct Mesh
{
	Vertex[] vtx;
	u32[]    idx;
	u32      offset;
	u32      length;
	u32      mat_id;
}

/**************************
 ***** GFX STRUCTURES *****
 **************************/

__gshared string g_BASE_ASSETS_DIR;

static u32
MagicValue(string str)
{
	assert(str.length == 4, "Magic value must 4 characters");
	return cast(u32)(cast(u32)(str[0] << 24) | cast(u32)(str[1] << 16) | cast(u32)(str[2] << 8) | cast(u32)(str[3] << 0));
}

void
SetBaseAssetsDir(string dir)
{
	if(!isDir(dir))
	{
		assert(false, "Unable to find assets directory");
	}

	g_BASE_ASSETS_DIR = dir;
}

void
U32ColToVec4(u32 col, Vec4* vec)
{
	if(!col)
	{
		vec.rgb = 0.0;
		vec.a   = 1.0;
	}
	else
	{
		vec.r = cast(f32)((col >>  0)&0xFF) / 255;
		vec.g = cast(f32)((col >>  8)&0xFF) / 255;
		vec.b = cast(f32)((col >> 16)&0xFF) / 255;
		vec.a = cast(f32)((col >> 24)&0xFF) / 255;
	}
}

// TODO: make better
u8[]
OpenFile(string file_name)
{
	File f;
	u8[] data;

	try
	{
		f = File(file_name, "rb");
		data = Alloc!(u8)(f.size());
		f.rawRead(data);
	}
	catch(Exception e)
	{
		data = null;
	}

	f.close();

	return data;
}

extern (C) cgltf_result
GLTFLoadCallback(cgltf_memory_options* memory_opts, cgltf_file_options* file_opts, const(char)* path, cgltf_size* size, void** data)
{
	u8[] file_data = OpenFile(ConvToStr(path[0 .. strlen(path)]));

	if(file_data == null) return cgltf_result_io_error;

	*size = cast(cgltf_size)file_data.length;
	*data = file_data.ptr;

	return cgltf_result_success;
}

extern (C) void
GLTFFreeCallback(cgltf_memory_options* memory_opts, cgltf_file_options* file_opts, void* data, cgltf_size size)
{
	Free(data);
}

ImageData
LoadImage(void* data, i32 size)
{
	ImageData image;

	i32 w, h, has_ch;
	u8* image_bytes = stbi_load_from_memory(cast(const(u8)*)data, size, &w, &h, &has_ch, 4);
	if(w <= 0 || h <= 0 || has_ch <= 0)
	{
		Logf("Failed to load image data");
	}
	else
	{
		image.data = image_bytes[0 .. w*h*has_ch];
		image.w    = w;
		image.h    = h;
		image.ch   = has_ch;
	}

	return image;
}

void
UnloadImage(ImageData* image)
{
	stbi_image_free(image.data.ptr);
}

ImageData
LoadImage(cgltf_image* asset_image, string tex_path)
{
	ImageData image;

	if(asset_image)
	{
		if(asset_image.uri != null)
		{
			if(strlen(asset_image.uri) > 5 && ConvToStr(asset_image.uri[0 .. 5]) == "data:")
			{
				u32 i;
				for(; asset_image.uri[i] != ',' && asset_image.uri[i] != 0; i += 1)	{}

				if(asset_image.uri[i] == 0) 
				{
					Logf("gltf data URI for is not a valid image");
				}
				else
				{
					u64 base_64_size = strlen(asset_image.uri + i + 1);
					for(; asset_image.uri[i + base_64_size] == '='; base_64_size -= 1) {}

					u64 bit_count = base_64_size*6 - (base_64_size*6) % 8;
					i32 out_size = cast(i32)(bit_count/8);

					void* data;
					cgltf_options options;
					options.file.read    = &GLTFLoadCallback;
					options.file.release = &GLTFFreeCallback;

					cgltf_result result = cgltf_load_buffer_base64(&options, out_size, asset_image.uri + i + 1, &data);
					if(result == cgltf_result_success)
					{
						image = LoadImage(data, out_size);
						cgltf_free(cast(cgltf_data*)data);
					}
				}
			}
			else
			{
				char[512] buf;
				char[] file_path = buf.sformat("%s%s%s", tex_path, dirSeparator, cast(char*)asset_image.uri);
				u8[] image_file = OpenFile(ConvToStr(file_path));

				image = LoadImage(image_file.ptr, cast(i32)image_file.length);

				Free(image_file);
			}
		}
		else if(asset_image.buffer_view != null && asset_image.buffer_view.buffer.data != null)
		{
			u8[] data = Alloc!(u8)(asset_image.buffer_view.size);

			i32 offset = cast(i32)(asset_image.buffer_view.offset);
			i32 stride = cast(i32)(asset_image.buffer_view.stride ? asset_image.buffer_view.stride : 1);

			u64 len = asset_image.buffer_view.size - (asset_image.buffer_view.size % stride);
			data[0 .. len] = (cast(u8*)(asset_image.buffer_view.buffer.data))[offset .. offset+len];

			string mime_type = ConvToStr(asset_image.mime_type[0 .. strlen(asset_image.mime_type)]);

			const string[] accepted_types = ["image\\/png", "image/png", "image\\/jpeg", "image/jpeg", "image\\/tga", "image/tga", "image\\/bmp", "image/bmp"];
			bool accepted;
			static foreach(type; accepted_types)
			{
				if(type == mime_type)
				{
					accepted = true;
					goto PostTypeCheck;
				}
			}

			PostTypeCheck:
			u8[] image_data;
			if(accepted)
			{
				image = LoadImage(data.ptr, cast(i32)data.length);
			}
			else 
			{
				Logf("Unknown image type [%s], unable to load", mime_type);
			}

			Free(data);
		}
	}

	if(image.data == null)
	{
		Logf("Failed to load GLTF image data/file");
		image = DEFAULT_IMAGE;
	}

	return image;
}

ModelData
LoadGLTF(Arena* arena, string file_name)
{
	ModelData model;

	u8[] file_data = OpenFile(file_name);
	assert(file_data != null);

	cgltf_options opts;
	cgltf_data*   data;

	opts.file.read    = &GLTFLoadCallback;
	opts.file.release = &GLTFFreeCallback;

	cgltf_result result = cgltf_parse(&opts, file_data.ptr, file_data.length, &data);

	if(result == cgltf_result_success)
	{
		result = cgltf_load_buffers(&opts, data, file_name.ptr);
		if(result != cgltf_result_success)
		{
			Logf("%s Failure: Unable to load buffers", __FUNCTION__);
		}

		u64 primitive_count;
		u64 vtx_count;
		u64 idx_count;

		for(u64 i = 0; i < data.nodes_count; i += 1)
		{
			cgltf_node* node = &data.nodes[i];

			if(node.mesh == null) continue;

			for(u64 j = 0; j < node.mesh.primitives_count; j += 1)
			{
				if(node.mesh.primitives[j].type == cgltf_primitive_type_triangles)
				{
					primitive_count += 1;

					if(node.mesh.primitives[j].indices != null && node.mesh.primitives[j].indices.buffer_view != null)
					{
						idx_count += node.mesh.primitives[j].indices.count;
					}

					for(u64 k = 0; k < node.mesh.primitives[j].attributes_count; k += 1)
					{
						if(node.mesh.primitives[j].attributes[k].type == cgltf_attribute_type_position)
						{
							vtx_count += node.mesh.primitives[j].attributes[k].data.count;
						}


					}
				}
			}
		}

		model.idx_buf = Alloc!(u32)(idx_count);
		model.vtx_buf = Alloc!(Vertex)(vtx_count);
		model.meshes  = Alloc!(Mesh)(primitive_count);
		model.mats    = Alloc!(Material)(data.materials_count+1);
		model.tex     = Alloc!(ImageData)(data.textures_count+1);

		model.mats[0] = DEFAULT_MATERIAL;

		u32 tex_idx = 1;
		string file_path = GetFilePath(file_name);
		for(u64 i = 0; i < data.materials_count; i += 1)
		{
			u64 mat_idx = i+1;

			model.mats[mat_idx] = DEFAULT_MATERIAL;

			if(data.materials[i].has_pbr_metallic_roughness)
			{
				auto pbr_mr = &data.materials[i].pbr_metallic_roughness;
				if(pbr_mr.base_color_texture.texture)
				{
					model.tex[tex_idx] = LoadImage(pbr_mr.base_color_texture.texture.image, file_path);
					model.mats[mat_idx].maps[MMI.Albedo].tex_id = tex_idx++;
				}

				model.mats[mat_idx].maps[MMI.Albedo].col.v = pbr_mr.base_color_factor;

				if(pbr_mr.metallic_roughness_texture.texture)
				{
					model.tex[tex_idx] = LoadImage(pbr_mr.base_color_texture.texture.image, file_path);
					model.mats[mat_idx].maps[MMI.Metallic].tex_id  = tex_idx;
					model.mats[mat_idx].maps[MMI.Roughness].tex_id = tex_idx;
					tex_idx += 1;

					model.mats[mat_idx].maps[MMI.Metallic].value  = pbr_mr.metallic_factor;
					model.mats[mat_idx].maps[MMI.Roughness].value = pbr_mr.roughness_factor;
				}

				if(data.materials[i].normal_texture.texture)
				{
					model.tex[tex_idx] = LoadImage(data.materials[i].normal_texture.texture.image, file_path);
					model.mats[mat_idx].maps[MMI.Normal].tex_id = tex_idx++;
				}

				if(data.materials[i].occlusion_texture.texture)
				{
					model.tex[tex_idx] = LoadImage(data.materials[i].occlusion_texture.texture.image, file_path);
					model.mats[mat_idx].maps[MMI.Occlusion].tex_id = tex_idx++;
				}

				if(data.materials[i].emissive_texture.texture)
				{
					model.tex[tex_idx] = LoadImage(data.materials[i].emissive_texture.texture.image, file_path);
					model.mats[mat_idx].maps[MMI.Emission].tex_id = tex_idx++;

					model.mats[mat_idx].maps[MMI.Emission].col.v[0 .. 3] = data.materials[i].emissive_factor;
					model.mats[mat_idx].maps[MMI.Emission].col.a         = 1.0;
				}
			}
		}

		u64 mesh_index, idx_index;
		vtx_count = 0;
		for(u64 i = 0; i < data.nodes_count; i += 1)
		{
			cgltf_node* node = &data.nodes[i];
			cgltf_mesh* mesh = node.mesh;

			if(mesh == null) continue;

			cgltf_float[16] world_transform;
			cgltf_node_transform_world(node, world_transform.ptr);

			Mat4 world_matrix = Mat4(world_transform);

			for(u64 p = 0; p < mesh.primitives_count; p += 1)
			{
				if(mesh.primitives[p].type != cgltf_primitive_type_triangles)
				{
					Logf("Unable to process primitive type [%s]", mesh.primitives[p].type);
					continue;
				}

				auto prim = mesh.primitives + p;	
				for(u64 j = 0; j < prim.attributes_count; j += 1)
				{
					switch(prim.attributes[j].type)
					{
						case cgltf_attribute_type_position:
							{
								cgltf_accessor* attr = prim.attributes[j].data;

								if(attr.type == cgltf_type_vec3 && attr.component_type == cgltf_component_type_r_32f)
								{
									AddMeshVertices(attr, &model, mesh_index, &vtx_count);

									f32* buffer = GetGLTFBuffer!(f32)(attr);
									for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
									{
										model.meshes[mesh_index].vtx[k].pos = Vec3(buffer[3*k+0], buffer[3*k+1], buffer[3*k+2]);
										Transform(&model.meshes[mesh_index].vtx[k].pos, &world_matrix, 1.0);
									}
								}
								else Logf("Format needs to be Vec3 f32, skipping.");
							} break;
						case cgltf_attribute_type_normal:
							{
								cgltf_accessor* attr = prim.attributes[j].data;

								if(attr.type == cgltf_type_vec3 && attr.component_type == cgltf_component_type_r_32f)
								{
									AddMeshVertices(attr, &model, mesh_index, &vtx_count);

									f32* buffer = GetGLTFBuffer!(f32)(attr);
									for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
									{
										model.meshes[mesh_index].vtx[k].normal = Vec3(buffer[3*k+0], buffer[3*k+1], buffer[3*k+2]);
										Transform(&model.meshes[mesh_index].vtx[k].normal, &world_matrix, 1.0);
									}
								}
								else Logf("Format needs to be Vec3 f32, skipping.");
							} break;
						case cgltf_attribute_type_tangent:
							{
								cgltf_accessor* attr = prim.attributes[j].data;

								if(attr.type == cgltf_type_vec4 && attr.component_type == cgltf_component_type_r_32f)
								{
									AddMeshVertices(attr, &model, mesh_index, &vtx_count);

									f32* buffer = GetGLTFBuffer!(f32)(attr);
									for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
									{
										model.meshes[mesh_index].vtx[k].tangent = world_matrix * Vec4(buffer[4*k+0], buffer[4*k+1], buffer[4*k+2], buffer[4*k+3]);
									}
								}
								else Logf("Format needs to be Vec4 f32, skipping.");
							} break;
						case cgltf_attribute_type_texcoord:
							{
								cgltf_accessor* attr = prim.attributes[j].data;

								if(attr.type == cgltf_type_vec2)
								{
									AddMeshVertices(attr, &model, mesh_index, &vtx_count);

									u32 uv_idx = prim.attributes[j].index;
									if(uv_idx >= 2)
									{
										Logf("Unable to use more than two uvs, skipping");
										continue;
									}

									if(attr.component_type == cgltf_component_type_r_32f)
									{
										f32* buffer = GetGLTFBuffer!(f32)(attr);
										for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
										{
											model.meshes[mesh_index].vtx[k].uv[uv_idx] = Vec2(buffer[2*k+0], buffer[2*k+1]);
										}
									}
									else if(attr.component_type == cgltf_component_type_r_8u)
									{
										u8* buffer = GetGLTFBuffer!(u8)(attr);
										for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
										{
											model.meshes[mesh_index].vtx[k].uv[uv_idx] = Vec2(cast(f32)(buffer[2*k+0])/255.0, cast(f32)(buffer[2*k+1])/255.0);
										}
									}
									else if(attr.component_type == cgltf_component_type_r_16u)
									{
										u16* buffer = GetGLTFBuffer!(u16)(attr);
										for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
										{
											model.meshes[mesh_index].vtx[k].uv[uv_idx] = Vec2(cast(f32)(buffer[2*k+0])/65535.0, cast(f32)(buffer[2*k+1])/65535.0);
										}
									}
									else
									{
										Logf("Unsupported format [%s] for uv, skipping", attr.component_type);
									}
								}
							} break;
						case cgltf_attribute_type_color:
							{	
								cgltf_accessor* attr = prim.attributes[j].data;

								if(attr.type == cgltf_type_vec3 || attr.type == cgltf_type_vec4)
								{
									AddMeshVertices(attr, &model, mesh_index, &vtx_count);

									if(attr.type == cgltf_type_vec3)
									{
										if(attr.component_type == cgltf_component_type_r_8u)
										{
											u8* buffer = GetGLTFBuffer!(u8)(attr);
											for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
											{
												model.meshes[mesh_index].vtx[k].color = Vec4(
													cast(f32)(buffer[k*3+0])/255.0, 
													cast(f32)(buffer[k*3+1])/255.0, 
													cast(f32)(buffer[k*3+2])/255.0,
													1.0
												);
											}
										}
										else if(attr.component_type == cgltf_component_type_r_16u)
										{
											u16* buffer = GetGLTFBuffer!(u16)(attr);
											for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
											{
												model.meshes[mesh_index].vtx[k].color = Vec4(
													cast(f32)(buffer[k*3+0])/65535.0,
													cast(f32)(buffer[k*3+1])/65535.0,
													cast(f32)(buffer[k*3+2])/65535.0,
													1.0
												);
											}
										}
										else if(attr.component_type == cgltf_component_type_r_32f)
										{
											f32* buffer = GetGLTFBuffer!(f32)(attr);
											for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
											{
												model.meshes[mesh_index].vtx[k].color = Vec4(buffer[k*3+0], buffer[k*3+1], buffer[k*3+2], 1.0);
											}
										}
										else
										{
											Logf("Color component type [%s] not supported", attr.component_type);
										}
									}
									else if(attr.type == cgltf_type_vec4)
									{
										if(attr.component_type == cgltf_component_type_r_8u)
										{
											u8* buffer = GetGLTFBuffer!(u8)(attr);
											for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
											{
												model.meshes[mesh_index].vtx[k].color = Vec4(
													cast(f32)(buffer[k*4+0])/255.0, 
													cast(f32)(buffer[k*4+1])/255.0, 
													cast(f32)(buffer[k*4+2])/255.0,
													cast(f32)(buffer[k*4+3])/255.0
												);
											}
										}
										else if(attr.component_type == cgltf_component_type_r_16u)
										{
											u16* buffer = GetGLTFBuffer!(u16)(attr);
											for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
											{
												model.meshes[mesh_index].vtx[k].color = Vec4(
													cast(f32)(buffer[k*4+0])/65535.0,
													cast(f32)(buffer[k*4+1])/65535.0,
													cast(f32)(buffer[k*4+2])/65535.0,
													cast(f32)(buffer[k*4+3])/65535.0
												);
											}
										}
										else if(attr.component_type == cgltf_component_type_r_32f)
										{
											f32* buffer = GetGLTFBuffer!(f32)(attr);
											for(u64 k = 0; k < model.meshes[mesh_index].vtx.length; k += 1)
											{
												model.meshes[mesh_index].vtx[k].color = Vec4(buffer[k*4+0], buffer[k*4+1], buffer[k*4+2], buffer[k*4+3]);
											}
										}
										else
										{
											Logf("Color component type [%s] not supported", attr.component_type);
										}
										
									}
								}
							} break;
						default: break;
					}


				}

				if(prim.indices != null && prim.indices.buffer_view != null)
				{
					cgltf_accessor* attr = prim.indices;

					model.meshes[mesh_index].idx = model.idx_buf[idx_index .. idx_index+attr.count];
					idx_index += attr.count;
					
					if(attr.component_type == cgltf_component_type_r_16u)
					{
						u16* buffer = GetGLTFBuffer!(u16)(attr);
						for(u64 k = 0; k < model.meshes[mesh_index].idx.length; k += 1)
						{
							model.meshes[mesh_index].idx[k] = cast(u32)buffer[k];
						}
					}
					else if(attr.component_type == cgltf_component_type_r_8u)
					{
						u8* buffer = GetGLTFBuffer!(u8)(attr);
						for(u64 k = 0; k < model.meshes[mesh_index].idx.length; k += 1)
						{
							model.meshes[mesh_index].idx[k] = cast(u32)buffer[k];
						}
					}
					else if(attr.component_type == cgltf_component_type_r_32u)
					{
						u32* buffer = GetGLTFBuffer!(u32)(attr);
						for(u64 k = 0; k < model.meshes[mesh_index].idx.length; k += 1)
						{
							model.meshes[mesh_index].idx[k] = buffer[k];
						}
					}
					else Logf("Unsupported index type [%s]", attr.component_type);
				}
				else // unindexed mesh (may implement later)
				{
					Logf("Unindexed mesh encountered");
				}

				for(u32 k = 0; k < data.materials_count; k += 1)
				{
					if(data.materials + k == prim.material)
					{
						model.meshes[mesh_index].mat_id = k+1;
						break;
					}
				}

				mesh_index += 1;
			}
		}
	}

	return model;
}

void
AddMeshVertices(cgltf_accessor* accessor, ModelData* model, u64 mesh_index, u64* vtx_count)
{
	if(model.meshes[mesh_index].vtx == null)
	{
		model.meshes[mesh_index].vtx    = model.vtx_buf[*vtx_count .. *vtx_count+accessor.count];
		model.meshes[mesh_index].offset = cast(u32)*vtx_count;
		model.meshes[mesh_index].length = cast(u32)accessor.count;
		*vtx_count += accessor.count;
	}
}

T*
GetGLTFBuffer(T)(cgltf_accessor* accessor)
{
	return cast(T*)(accessor.buffer_view.buffer.data + accessor.buffer_view.offset/T.sizeof + accessor.offset/T.sizeof);
}

static ImageData
GenerateDefaultTexture(u32 x, u32 y)
{
	ImageData data = {
		w:  x,
		h:  y,
		ch: 4,
	};

	u64 tex_size = x*y*4;
	u8[] placeholder_tex = new u8[tex_size];

	u8[4] magenta = [255, 0, 255, 255];
	u8[4] black   = [0,   0, 0,   255];
	u64 half      = tex_size/2;
	for(u64 i = 0; i < tex_size; i += 32)
	{
		bool swap = i <= half;
		for(u64 j = 0; j < 16; j += 4)
		{
			placeholder_tex[i+j .. i+j+4]       = !swap ? magenta[0 .. $] : black[0 .. $];
			placeholder_tex[i+j+16 .. i+j+16+4] = !swap ? black[0 .. $]   : magenta[0 .. $];
		}
	}

	data.data = placeholder_tex;

	return data;
}

static Material
GenerateDefaultMaterial()
{
	Material mat;

	mat.maps[MMI.Albedo].col   = Vec4(1.0);
	mat.maps[MMI.Metallic].col = Vec4(1.0);

	return mat;
}