import aliases;
import includes;
import vulkan : Destroy;
import vulkan;
import assets;
import util;
import alloc;
import platform;
import math;
import core.stdc.math : cosf, sinf;
import std.algorithm.sorting;
import fonts;

f32 g_DELTA;

struct UIVertex
{
	Vec2 dst_start;
	Vec2 dst_end;
	Vec2 src_start;
	Vec2 src_end;
	Vec4 col;
}

struct UIPushConst
{
	Vec2 res;
}

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

struct MeshPart
{
	u32 mat;
	u32 offset;
	u32 length;
}

struct Model
{
	Vec3				pos = Vec3(0.0);
	Buffer 			vertex_buffer;
	Buffer 			index_buffer;
	MeshPart[] 	parts;

	string			name;

	Buffer[] 		materials;
	ImageView[] textures;

	Vec3[]			positions;
	u32[]				pos_indices;
	u32[]				indices;
}

struct Game
{
	Renderer rd;
	Arena 	 arena;
	Arena 	 frame_arena;

	PlatformWindow* window;

	Pipeline ui_pipeline;
	DescSetLayout ui_desc_layout;
	DescSet ui_desc_set;
	PipelineLayout ui_layout;

	UIVertex[] ui_vertex_buf;
	u32[]			 ui_index_buf;
	u32				 ui_count;

	ImageView 	default_tex;
	ImageView 	font_tex;

	UIPushConst ui_pc;

	Timer timer;
}

Game
InitGame(PlatformWindow* window)
{
	Game g = { 
		rd: InitRenderer(window, MB(24), MB(32)), 
		arena: CreateArena(MB(32)),
		frame_arena: CreateArena(MB(16)),
		window: window,
		timer: CreateTimer(),
	};

	UVec2 ext = GetExtent(&g.rd);

	DescLayoutBinding[] layout_bindings = [
		{ binding: 0, descriptorType: VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, descriptorCount: 1, stageFlags: VK_SHADER_STAGE_ALL },
	];

	g.ui_desc_layout = CreateDescSetLayout(&g.rd, layout_bindings);
	g.ui_desc_set = AllocDescSet(&g.rd, g.ui_desc_layout);
	g.ui_layout = CreatePipelineLayout(&g.rd, g.ui_desc_layout, UIPushConst.sizeof);

	u8[16*16*4] white_tex;
	white_tex[] = u8.max;

	u8[] atlas = LoadAssetData(&g.frame_arena, "textures/atlas.png");
	int width, height, has_ch;
	auto img = stbi_load_from_memory(atlas.ptr, cast(int)atlas.length, &width, &height, &has_ch, 4);
	assert(width == FONT_ATLAS.width && height == FONT_ATLAS.height && has_ch == 1, "atlas height and width do not match");
	u8[] img_slice = img[0 .. width * height * 4];

	CreateImageView(&g.rd, &g.font_tex, FONT_ATLAS.width, FONT_ATLAS.height, 4, img_slice);
	CreateImageView(&g.rd, &g.default_tex, 16, 16, 4, white_tex);

	WriteGUI(&g.rd, g.ui_desc_set, &g.font_tex);

	GfxPipelineInfo ui_info = {
		vertex_shader: "shaders/gui.vert.spv",
		frag_shader: "shaders/gui.frag.spv",
		input_rate: IR.Instance,
		input_rate_stride: UIVertex.sizeof,
		layout: g.ui_layout,
		vertex_attributes: [
			{ binding: 0, location: 0, format: FMT.RG_F32, offset: UIVertex.dst_start.offsetof },
			{ binding: 0, location: 1, format: FMT.RG_F32, offset: UIVertex.dst_end.offsetof },
			{ binding: 0, location: 2, format: FMT.RG_F32, offset: UIVertex.src_start.offsetof },
			{ binding: 0, location: 3, format: FMT.RG_F32, offset: UIVertex.src_end.offsetof },
			{ binding: 0, location: 4, format: FMT.RGBA_F32, offset: UIVertex.col.offsetof },
		],
	};

	g.ui_pipeline = CreateGraphicsPipeline(&g.rd, &ui_info);

	PrintShaderDisassembly(&g.rd, g.ui_pipeline, VK_SHADER_STAGE_VERTEX_BIT);
	PrintShaderDisassembly(&g.rd, g.ui_pipeline, VK_SHADER_STAGE_FRAGMENT_BIT);

	g.ui_vertex_buf = GetUIVertexBuffer(&g.rd);
	g.ui_index_buf = GetUIIndexBuffer(&g.rd);

	WaitForTransfers(&g.rd);

	Reset(&g.frame_arena);

	return g;
}

void
Copy(Model* dst, Model* src)
{
	dst.vertex_buffer = src.vertex_buffer;
	dst.index_buffer = src.index_buffer;
	dst.parts = src.parts;
	dst.materials = src.materials;
	dst.textures = src.textures;
}

void
ProcessInputs(Game* g, Camera* cam)
{
	HandleInputs(cam, &g.window.inputs);
}

void
Cycle(Game* g)
{
	g_DELTA = DeltaTime(&g.timer);

	g.ui_count = 0;

	Reset(&g.frame_arena);

	DrawRect(g, 500.0, 500.0, 800.0, 800.0, Vec4(0.2, 0.3, 0.7, 1.0));
	DrawText(g, 200.0, 200.0, 32.0, "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ");

	BeginFrame(&g.rd);
	
	UVec2 ext = GetExtent(&g.rd);
	g.ui_pc.res.x = cast(f32)ext.x;
	g.ui_pc.res.y = cast(f32)ext.y;
	
	BeginRendering(&g.rd);

	PushConstants(&g.rd, g.ui_pipeline, &g.ui_pc);

	Bind(&g.rd, g.ui_pipeline, g.ui_desc_set);

	BindUIBuffers(&g.rd);

	DrawIndexed(&g.rd, 6, g.ui_count, 0);

	FinishRendering(&g.rd);

	SubmitAndPresent(&g.rd);
}

pragma(inline): void
DrawModel(Game* g, Model* model)
{
	BindBuffers(&g.rd, &model.index_buffer, &model.vertex_buffer);
	foreach(i, part; model.parts)
	{
		DrawIndexed(&g.rd, part.length, 1, part.offset);
	}
}

void
Destroy(Game* g)
{
	WaitIdle(&g.rd);

	Destroy(&g.rd, g.ui_pipeline);
	Destroy(&g.rd);
}

void
Destroy(Renderer* rd, Model* model)
{
	Destroy(rd, &model.vertex_buffer);
	Destroy(rd, &model.index_buffer);
	
	foreach(i, mat; model.materials)
	{
		Destroy(rd, model.materials.ptr + i);
	}

	foreach(i, view; model.textures)
	{
		Destroy(rd, model.textures.ptr + i);
	}
}

struct Camera
{
	Vec3 velocity = Vec3(0.0);
	Vec3 pos = Vec3(0.0, 0.0, 2.0);
	f32 pitch = 0.0;
	f32 yaw = 0.0;
}

pragma(inline): void
HandleInputs(Camera* cam, Inputs* inputs)
{
	foreach(i; 0 .. inputs.count)
	{
		InputEvent event = inputs.events[i];
		switch(event.key)
		{
			case Input.W: cam.velocity.z = event.pressed ? -1.0 : 0.0; break;
			case Input.S: cam.velocity.z = event.pressed ? 1.0 : 0.0; break;
			case Input.A: cam.velocity.x = event.pressed ? -1.0 : 0.0; break;
			case Input.D: cam.velocity.x = event.pressed ? 1.0 : 0.0; break;
			case Input.Space: cam.velocity.y = event.pressed ? 1.0 : 0.0; break;
			case Input.LeftCtrl: cam.velocity.y = event.pressed ? -1.0 : 0.0; break;
			case Input.MouseMotion:
			 	{
				 	// (0, 0) top left
					cam.yaw += cast(f32)(event.rel_x) / 200.0;
					cam.pitch += cast(f32)(event.rel_y) / 200.0;
			 } break;
			default: break;
		}
	}
}

Mat4
RotationMatrix(Camera* cam)
{
	Quat pitch = QuatFromAxis(cam.pitch, Vec3(1.0, 0.0, 0.0));
	Quat yaw = QuatFromAxis(cam.yaw, Vec3(0.0, 1.0, 0.0));
	return cast(Mat4)(yaw) * cast(Mat4)(pitch);
}

Mat4
ViewMatrix(Camera* cam)
{
	Mat4 translation = Mat4Identity();
	Translate(&translation, cam.pos);
	Mat4 rotation = RotationMatrix(cam);
	return Inverse(translation * rotation);
}

void
Update(Camera* cam)
{
	Mat4 rotation = RotationMatrix(cam);
	Vec4 pos = rotation * Vec4(cam.velocity * g_DELTA, 0.0);
	cam.pos += pos.xyz;
}

void
Sort(Game* g, Vec3 pos, Model* model)
{
	f32[] lengths = AllocArray!(f32)(&g.frame_arena, model.positions.length);

	foreach(i; 0 .. lengths.length)
	{
		model.pos_indices[i] = cast(u32)i;
		lengths.ptr[i] = fabs(Norm(&pos) - Norm(model.positions.ptr + i));
	}

	makeIndex!("a < b")(lengths, model.pos_indices);

	Logf("%s", model.positions.length);

	foreach(i, v; model.pos_indices)
	{
		model.indices[i+0] = cast(u32)((3*v)+0);
		model.indices[i+1] = cast(u32)((3*v)+1);
		model.indices[i+2] = cast(u32)((3*v)+2);
	}
	
	UpdateIndexBuffer(&g.rd, model);
}

void
DrawText(Game* g, f32 x, f32 y, f32 px, string str)
{
	f32 x_pos = x;
	foreach(ch; str)
	{
		foreach(glyph; FONT_ATLAS.glyphs)
		{
			if (ch == glyph.ch)
			{
				UIVertex* v = g.ui_vertex_buf.ptr + g.ui_count;

				f32 w = px * (glyph.plane_right - glyph.plane_left);
				f32 h = px * (glyph.plane_bottom - glyph.plane_top);
				f32 y_pos = px * glyph.plane_bottom;

				v.dst_start.x = x_pos;
				v.dst_start.y = y + h - y_pos;
				v.dst_end.x = x_pos + w;
				v.dst_end.y = y - y_pos;

				v.src_start.x = glyph.atlas_left;
				v.src_start.y = glyph.atlas_top;
				v.src_end.x = glyph.atlas_right;
				v.src_end.y = glyph.atlas_bottom;

				v.col = Vec4(1.0, 1.0, 1.0, 1.0);

				x_pos += px * glyph.advance;

				AddUIIndices(g);
			}
		}
	}
}

void
AddUIIndices(Game* g)
{
	g.ui_index_buf[0] = 0; 
	g.ui_index_buf[1] = 1; 
	g.ui_index_buf[2] = 2; 
	g.ui_index_buf[3] = 2; 
	g.ui_index_buf[4] = 1; 
	g.ui_index_buf[5] = 3; 

	g.ui_count += 1;
}

void
DrawRect(Game* g, f32 p0_x, f32 p0_y, f32 p1_x, f32 p1_y, Vec4 col)
{
	// Y reversed
	g.ui_vertex_buf[g.ui_count].dst_start.x = p0_x;
	g.ui_vertex_buf[g.ui_count].dst_start.y = p0_y;
	g.ui_vertex_buf[g.ui_count].dst_end.x = p1_x;
	g.ui_vertex_buf[g.ui_count].dst_end.y = p1_y;
	g.ui_vertex_buf[g.ui_count].col = col;

	AddUIIndices(g);
}

// TODO: integrate this with vulkan again
Model
LoadModel(Game* g, string name)
{
	AssetInfo info; // = GetAssetInfo(name);

	u8[] data = LoadAssetData(&g.frame_arena, name);
	Model model = {
		name: name,
	};

	m3d_t* m3d = m3d_load(data.ptr, null, null, null);
	scope(exit) m3d_free(m3d);

	u32[] tex_lookup = AllocArray!(u32)(&g.frame_arena, m3d.numtexture);
	model.textures = AllocArray!(ImageView)(&g.arena, m3d.numtexture);
	foreach(i; 0 .. m3d.numtexture)
	{
		u32 w = m3d.texture[i].w; u32 h = m3d.texture[i].h; u32 ch = m3d.texture[i].f;
		u8[] tex_data = m3d.texture[i].d[0 .. w * h * ch];

		const(char)[] tex_name = m3d.texture[i].name[0 .. strlen(m3d.texture[i].name)];
		CreateImageView(&g.rd, &model.textures[i], w, h, ch, tex_data);

		//tex_lookup[i] = Pop(&g.rd, DT.SampledImage);
	}

	Material[] mats = AllocArray!(Material)(&g.frame_arena, m3d.nummaterial);
	u32[] mat_lookup = AllocArray!(u32)(&g.frame_arena, m3d.nummaterial);
	model.materials = AllocArray!(Buffer)(&g.arena, m3d.nummaterial);
	foreach(i; 0 .. m3d.nummaterial)
	{
		const(char)[] mat_name = m3d.material[i].name[0 .. strlen(m3d.material[i].name)];

		foreach(j; 0 .. m3d.material[i].numprop)
		{
			switch (m3d.material[i].prop[j].type)
			{
			case m3dp_Kd: ConvertColor(&mats[i].albedo, m3d.material[i].prop[j].value.color); break;
			case m3dp_Ka: ConvertColor(&mats[i].ambient, m3d.material[i].prop[j].value.color); break;
			case m3dp_Ks: ConvertColor(&mats[i].specular, m3d.material[i].prop[j].value.color); break;
			case m3dp_Ns: mats[i].shininess = m3d.material[i].prop[j].value.fnum; break;
			case m3dp_d:  mats[i].alpha = m3d.material[i].prop[j].value.fnum; break;
			case m3dp_map_Kd:
				{
					mats[i].albedo_texture = tex_lookup[m3d.material[i].prop[j].value.textureid];
					mats[i].albedo_has_texture = true;
				} break;
			case m3dp_map_Ka:
				{
					mats[i].ambient_texture = tex_lookup[m3d.material[i].prop[j].value.textureid];
					mats[i].ambient_has_texture = true;
				} break;
			case m3dp_map_Ks:
				{
					mats[i].specular_texture = tex_lookup[m3d.material[i].prop[j].value.textureid];
					mats[i].specular_has_texture = true;
				} break;
			case m3dp_map_D:
				{
					mats[i].alpha_texture = tex_lookup[m3d.material[i].prop[j].value.textureid];
					mats[i].alpha_has_texture = true;
				} break;
			default: Logf("Unsupported property: %s", M3DPropToStr(m3d.material[i].prop[j].type)); break;
			}
		}

		CreateBuffer(&g.rd, &model.materials[i], BT.Uniform, Material.sizeof, false);
		assert(Transfer(&g.rd, &model.materials[i], &mats[i]), "LoadModel failure: Transfer error when transferring material");

		//mat_lookup[i] = Pop(&g.rd, DT.Material);
	}

	u32 mesh_count = 0;
	u64 last = u64.max;
	foreach(i; 0 .. m3d.numface)
	{
		if (m3d.face[i].materialid != last)
		{
			last = m3d.face[i].materialid;
			mesh_count += 1;
		}
	}

	model.parts = AllocArray!(MeshPart)(&g.arena, mesh_count);
	last = u64.max;
	u32 index = 0;
	foreach(i; 0 .. m3d.numface)
	{
		if (last == u64.max)
		{
			model.parts[index].mat = (m3d.face[i].materialid != u32.max) ? mat_lookup[m3d.face[i].materialid] : 0;
			model.parts[index].offset = 0;
			last = m3d.face[i].materialid;
		}
		else if (m3d.face[i].materialid != last)
		{
			u32 vertex_index = i * 3;

			model.parts[index].length = vertex_index - model.parts[index].offset;
			index += 1;

			model.parts[index].mat = mat_lookup[m3d.face[i].materialid];
			model.parts[index].offset = vertex_index;
			last = m3d.face[i].materialid;
		}
		else if (i == m3d.numface-1)
		{
			u32 vertex_index = i * 3;
			model.parts[index].length = vertex_index+3 - model.parts[index].offset;
		}
	}

	m3dv_t* vertex;

	u32[] indices = AllocArray!(u32)(&g.arena, m3d.numface * 3);
	Vertex[] vertices = AllocArray!(Vertex)(&g.frame_arena, m3d.numface * 3);
	Vec3[] positions = AllocArray!(Vec3)(&g.arena, m3d.numface);
	u32[] pos_indices = AllocArray!(u32)(&g.arena, m3d.numface);

	foreach(i; 0 .. m3d.numface)
	{
		u32 vi = (i * 3);
		u32 i0 = vi+0;
		u32 i1 = vi+1;
		u32 i2 = vi+2;

		CopyVertex(&vertices[i0].pos, &m3d.vertex[m3d.face[i].vertex[0]]);
		CopyVertex(&vertices[i1].pos, &m3d.vertex[m3d.face[i].vertex[1]]);
		CopyVertex(&vertices[i2].pos, &m3d.vertex[m3d.face[i].vertex[2]]);

		CopyVertex(&vertices[i0].normal, &m3d.vertex[m3d.face[i].normal[0]]);
		CopyVertex(&vertices[i1].normal, &m3d.vertex[m3d.face[i].normal[1]]);
		CopyVertex(&vertices[i2].normal, &m3d.vertex[m3d.face[i].normal[2]]);

		ConvertColor(&vertices[i0].color, m3d.vertex[m3d.face[i].vertex[0]].color);
		ConvertColor(&vertices[i1].color, m3d.vertex[m3d.face[i].vertex[1]].color);
		ConvertColor(&vertices[i2].color, m3d.vertex[m3d.face[i].vertex[2]].color);

		if (m3d.numtmap)
		{
			vertices[i0].uv.x = m3d.tmap[m3d.face[i].texcoord[0]].u;
			vertices[i0].uv.y = m3d.tmap[m3d.face[i].texcoord[0]].v;

			vertices[i1].uv.x = m3d.tmap[m3d.face[i].texcoord[1]].u;
			vertices[i1].uv.y = m3d.tmap[m3d.face[i].texcoord[1]].v;

			vertices[i2].uv.x = m3d.tmap[m3d.face[i].texcoord[2]].u;
			vertices[i2].uv.y = m3d.tmap[m3d.face[i].texcoord[2]].v;
		}

		indices[i0] = i0;
		indices[i1] = i1;
		indices[i2] = i2;

		Vec3 center = (vertices[i0].pos + vertices[i1].pos + vertices[i2].pos) / 3.0;

		positions[i] = center;
		pos_indices[i] = i;
	}

	for(u64 i = 0; i < indices.length; i += 3)
	{
		u32 i0 = indices[i+0];
		u32 i1 = indices[i+1];
		u32 i2 = indices[i+2];

		Vec3 edge1 = vertices[i1].pos - vertices[i0].pos;
		Vec3 edge2 = vertices[i2].pos - vertices[i0].pos;

		Vec2 delta_uv1 = vertices[i1].uv - vertices[i0].uv;
		Vec2 delta_uv2 = vertices[i2].uv - vertices[i0].uv;

		f32 dividend = delta_uv1.x * delta_uv2.y - delta_uv2.x * delta_uv1.y;
		f32 fc = 1.0 / dividend;

		Vec3 tangent = Vec3(
				fc * (delta_uv2.y * edge1.x - delta_uv1.y * edge2.x),
				fc * (delta_uv2.y * edge1.y - delta_uv1.y * edge2.y),
				fc * (delta_uv2.y * edge1.z - delta_uv1.y * edge2.z)
		);

		Normalize(&tangent);

		f32 handedness = ((delta_uv1.y * delta_uv2.x - delta_uv2.y * delta_uv1.x) < 0.0) ? -1.0 : 1.0;

		Vec3 t = tangent * handedness;

		vertices[i0].tangent = Vec4(t, handedness);
		vertices[i1].tangent = Vec4(t, handedness);
		vertices[i2].tangent = Vec4(t, handedness);
	}

	assert(vertices.length > 0 && indices.length > 0, "Error loading model");

	CreateBuffer(&g.rd, &model.vertex_buffer, BT.Vertex, vertices.length * Vertex.sizeof, false);
	CreateBuffer(&g.rd, &model.index_buffer, BT.Index, indices.length * u32.sizeof, false);

	assert(Transfer(&g.rd, &model.vertex_buffer, vertices), "LoadModel failure: Unable to transfer vertex buffer");
	assert(Transfer(&g.rd, &model.index_buffer, indices), "LoadModel failure: Unable to transfer index buffer");

	//WriteDescriptors(&g.rd, DT.Material, model.materials, mat_lookup);
	//WriteDescriptors(&g.rd, DT.SampledImage, model.textures, tex_lookup);

	model.positions = positions;
	model.indices = indices;
	model.pos_indices = pos_indices;

	return model;
}

string
M3DPropToStr(u8 type)
{
	string result = "Unknown";
	switch(type)
	{
		case m3dp_Kd: result = "Diffuse"; break;
		case m3dp_Ka: result = "Ambient"; break;
		case m3dp_Ks: result = "Specular Color"; break;
		case m3dp_Ns: result = "Specular Exponent"; break;
		case m3dp_Ke: result = "Emissive"; break;
		case m3dp_Tf: result = "Transmission"; break;
		case m3dp_Km: result = "Bump Strength"; break;
		case m3dp_d: result = "Alpha"; break;
		case m3dp_il: result = "Illumination"; break;
		case m3dp_Pr: result = "Roughness"; break;
		case m3dp_Pm: result = "Metallic"; break;
		case m3dp_Ps: result = "Sheen"; break;
		case m3dp_Ni: result = "Refraction"; break;
		case m3dp_Nt: result = "Face Thickness"; break;
		case m3dp_map_Kd: result = "Diffuse Texture"; break;
		case m3dp_map_Ka: result = "Ambient Texture"; break;
		case m3dp_map_Ks: result = "Specular Texture"; break;
		case m3dp_map_Ns: result = "Specular Exponent Texture"; break;
		case m3dp_map_Ke: result = "Emissive Texture"; break;
		case m3dp_map_Tf: result = "Transmission Texture"; break;
		case m3dp_map_Km: result = "Bump Map"; break;
		case m3dp_map_D: result = "Alpha Map"; break;
		case m3dp_map_N: result = "Normal Map"; break;
		case m3dp_map_Pr: result = "Roughness Map"; break;
		case m3dp_map_Pm: result = "Metallic Map"; break;
		case m3dp_map_Ps: result = "Sheen Map"; break;
		case m3dp_map_Ni: result = "Refraction Map"; break;
		case m3dp_map_Nt: result = "Thickness Map"; break;
		default: break;
	}

	return result;
}

pragma(inline): void
CopyVertex(Vec3* dst, m3dv_t* src)
{
	dst.x = src.x;
	dst.y = src.y;
	dst.z = src.z;
}

void
ReadModel(Game* g, string name)
{
	AssetInfo info; // = GetAssetInfo(name);

	u8[] data = LoadAssetData(&g.frame_arena, name);

	m3d_t* m3d = m3d_load(data.ptr, null, null, null);

	foreach(i; 0 .. m3d.numtexture)
	{
		const(char)[] tex_name = m3d.texture[i].name[0 .. strlen(m3d.texture[i].name)];
	}

	if (m3d.numtexture == 0)
	{
		Log("No textures in model");
	}

	foreach(i; 0 .. m3d.nummaterial)
	{
		const(char)[] mat_name = m3d.material[i].name[0 .. strlen(m3d.material[i].name)];
	}
}