import aliases;
import includes;
import assets;
import util;
import alloc;
import vulkan;
import vulkan : Destroy, Init, Draw, DrawIndexed, Bind, BindUIBuffers, BeginRender, SetUniform, PrepCompute, Dispatch, FinishFrame, BeginFrame;
import assets;
import std.math.traits : isNaN;
import u = util;
import p = platform;

alias Shader = VkShaderModule;
alias Pipeline = PipelineHandle;
alias Attribute = VkVertexInputAttributeDescription;

enum InputRate : int
{
	Vertex = VK_VERTEX_INPUT_RATE_VERTEX,
	Instance = VK_VERTEX_INPUT_RATE_INSTANCE,
}

alias IR = InputRate;

enum Format : VkFormat
{
	UINT = VK_FORMAT_R32_UINT,
	R_F32 = VK_FORMAT_R32_SFLOAT,
	RG_F32 = VK_FORMAT_R32G32_SFLOAT,
	RGB_F32 = VK_FORMAT_R32G32B32_SFLOAT,
	RGBA_F32 = VK_FORMAT_R32G32B32A32_SFLOAT,
	RGBA_UINT = VK_FORMAT_B8G8R8A8_UINT,
	RGBA_UNORM = VK_FORMAT_R8G8B8A8_UNORM,
}

alias FMT = Format;

enum BufferType : int
{
	None = 0,
	Vertex = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
	Index = VK_BUFFER_USAGE_INDEX_BUFFER_BIT,
	Uniform = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
	Storage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
	Staging = VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
}

alias BT = BufferType;

enum PipelineType : int
{
	Graphics,
	Compute,
}

alias PT = PipelineType;

struct GfxPipelineInfo
{
	string vertex_shader;
	string frag_shader;
	InputRate input_rate;
	u32 input_rate_stride;
	Attribute[] vertex_attributes;
}

struct Renderer
{
	Arena arena;
	Arena temp_arena;

	Vulkan vk;
	p.Window* window;

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

	Pipeline pbr_pipeline;
	Pipeline triangle_pipeline;
	Pipeline compute_pipeline;
	Pipeline ui_pipeline;

	Model yoder;
}

struct GlobalUniforms
{
	Vec2 res;
}

struct ShaderUniforms
{
	f32 placeholder;
}

struct Material
{
	u32 albedo_texture;
	u32 ambient_texture;
	u32 specular_texture;
	b32 albedo_has_texture;
	b32 ambient_has_texture;
	b32 specular_has_texture;
	Vec4 ambient;
	Vec4 diffuse;
	Vec4 specular;
	f32 shininess = 0.0;
}

struct UIVertex
{
	Vec2 p0;
	Vec2 p1;
	Vec4 col;
	u32 texture;
}

struct Vertex
{
	Vec4 pos;
	Vec4 n;
	Vec2 uv;
	u32 col;
}

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

struct Model
{
	Buffer 			vertex_buffer;
	Buffer 			index_buffer;
	MeshPart[] 	parts;

	string			name;

	Buffer[] 		materials;
	ImageView[] textures;
}

Renderer
Init(p.Window* window)
{
	u.Result!(Vulkan) vk_result = Init(window, u.MB(24), u.MB(32));
	assert(vk_result.ok, "Init failure: Unable to initialize Vulkan");

	Renderer rd = {
		arena: CreateArena(u.MB(16)),
		temp_arena: CreateArena(u.MB(16)),
		vk: vk_result.value,
		window: window,
		ui_vertex_buf: GetUIVertexBuffer(&vk_result.value),
		ui_index_buf: GetUIIndexBuffer(&vk_result.value),
	};

	GfxPipelineInfo triangle_info = {
		vertex_shader: "shaders/triangle.vert",
		frag_shader: "shaders/triangle.frag",
	};

	GfxPipelineInfo ui_info = {
		vertex_shader: "shaders/gui.vert",
		frag_shader: "shaders/gui.frag",
		input_rate: IR.Instance,
		input_rate_stride: UIVertex.sizeof,
		vertex_attributes: [
			{ binding: 0, location: 0, format: FMT.RG_F32, offset: 0 },
			{ binding: 0, location: 1, format: FMT.RG_F32, offset: UIVertex.p1.offsetof },
			{ binding: 0, location: 2, format: FMT.RGBA_F32, offset: UIVertex.col.offsetof },
			{ binding: 0, location: 3, format: FMT.UINT, offset: UIVertex.texture.offsetof },
		],
	};

	GfxPipelineInfo pbr_info = {
		vertex_shader: "shaders/pbr.vert",
		frag_shader: "shaders/pbr.frag",
		input_rate_stride: Vertex.sizeof,
		vertex_attributes: [
			{ binding: 0, location: 0, format: FMT.RGBA_F32, offset: 0 },
			{ binding: 0, location: 1, format: FMT.RGBA_F32, offset: Vertex.n.offsetof },
			{ binding: 0, location: 2, format: FMT.RG_F32, offset: Vertex.uv.offsetof },
			{ binding: 0, location: 3, format: FMT.RGBA_UNORM, offset: Vertex.col.offsetof },
		],
	};

	rd.pbr_pipeline = BuildGfxPipeline(&rd, &pbr_info);
	rd.triangle_pipeline = BuildGfxPipeline(&rd, &triangle_info);
	rd.ui_pipeline = BuildGfxPipeline(&rd, &ui_info);
	rd.compute_pipeline = BuildCompPipeline(&rd, "shaders/gradient.comp");

	rd.yoder = LoadModel(&rd, "models/yoda");

	return rd;
}

bool
Cycle(Renderer* rd)
{
	rd.ui_count = 0;

	bool success = BeginFrame(rd);

	if (success)
	{
		Bind(rd, &rd.compute_pipeline);

		SetUniform(rd, &rd.globals);

		DrawRect(rd, 150.0, 300.0, 500.0, 700.0, Vec4(r: 0.0, g: 0.0, b: 1.0, a: 1.0));

		PrepCompute(rd);

		Dispatch(rd);

		BeginRender(rd);

		Bind(rd, &rd.ui_pipeline);

		BindUIBuffers(rd);

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

		Bind(rd, &rd.triangle_pipeline);

		Draw(rd, 3, 1);

		Bind(rd, &rd.pbr_pipeline);

		DrawModel(rd, &rd.yoder);

		success = FinishFrame(rd);
	}

	return success;
}

pragma(inline): void
DrawModel(Renderer* rd, Model* model)
{
	BindBuffers(&rd.vk, &model.index_buffer, &model.vertex_buffer);

	foreach(i, part; model.parts)
	{
		DrawIndexed(rd, part.length, 1, part.offset);
	}
}

pragma(inline): void
CopyVertex(Vec4* dst, m3dv_t* src)
{
	asm
	{
		mov R8, src;
		mov R9, dst;
		movups XMM0, src.x.offsetof[R8];
		movups dst.x.offsetof[R9], XMM0;
	}

debug
{
	assert(dst.x == src.x && dst.y == src.y && dst.z == src.z && dst.w == src.w, "Vertex copy failed");
	assert(!isNaN(dst.x) && !isNaN(dst.y) && !isNaN(dst.z) && !isNaN(dst.w), "Vertex contains NaN");
}
}

pragma(inline): bool
BeginFrame(Renderer* rd)
{
	return BeginFrame(&rd.vk);
}

pragma(inline): bool
FinishFrame(Renderer* rd)
{
	return FinishFrame(&rd.vk);
}

pragma(inline): void
Dispatch(Renderer* rd)
{
	Dispatch(&rd.vk);
}

pragma(inline): void
PrepCompute(Renderer* rd)
{
	PrepCompute(&rd.vk);
}

pragma(inline): void
SetUniform(Renderer* rd, GlobalUniforms* uniforms)
{
	SetUniform(&rd.vk, uniforms);
}

pragma(inline): void
BeginRender(Renderer* rd)
{
	BeginRender(&rd.vk);
}

pragma(inline): void
BindUIBuffers(Renderer* rd)
{
	BindUIBuffers(&rd.vk);
}

pragma(inline): void
DrawIndexed(Renderer* rd, u32 index_count, u32 instance_count, u32 index_offset)
{
	DrawIndexed(&rd.vk, index_count, instance_count, index_offset);
}

pragma(inline): void
Draw(Renderer* rd, u32 index_count, u32 instance_count)
{
	Draw(&rd.vk, index_count, instance_count);
}

pragma(inline): void 
Bind(Renderer* rd, Pipeline* pipeline)
{
	Bind(&rd.vk, pipeline);
}

void
DrawRect(Renderer* rd, f32 p0_x, f32 p0_y, f32 p1_x, f32 p1_y, Vec4 col)
{
	rd.ui_vertex_buf[rd.ui_count].p0.x = p0_x;
	rd.ui_vertex_buf[rd.ui_count].p0.y = p0_y;
	rd.ui_vertex_buf[rd.ui_count].p1.x = p1_x;
	rd.ui_vertex_buf[rd.ui_count].p1.y = p1_y;
	rd.ui_vertex_buf[rd.ui_count].col = col;

	u32 index_count = rd.ui_count * 6;

	rd.ui_index_buf[index_count+0] = index_count+0;
	rd.ui_index_buf[index_count+1] = index_count+1;
	rd.ui_index_buf[index_count+2] = index_count+2;
	rd.ui_index_buf[index_count+3] = index_count+1;
	rd.ui_index_buf[index_count+4] = index_count+2;
	rd.ui_index_buf[index_count+5] = index_count+3;

	rd.ui_count += 1;
}

pragma(inline): Pipeline
BuildGfxPipeline(Renderer* rd, GfxPipelineInfo* info)
{
	return CreateGraphicsPipeline(&rd.vk, info);
}

pragma(inline): Pipeline
BuildCompPipeline(Renderer* rd, string compute)
{
	return CreateComputePipeline(&rd.vk, compute);
}

Model
LoadModel(Renderer* rd, string name)
{
	Model model = {
		name: name,
	};

	u8[] data = LoadAssetData(&rd.temp_arena, name);
	m3d_t* m3d = m3d_load(data.ptr, null, null, null);

	u32[] tex_lookup = AllocArray!(u32)(&rd.temp_arena, m3d.numtexture);
	model.textures = AllocArray!(ImageView)(&rd.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;

		model.textures[i] = CreateImageView(&rd.vk, w, h, ch);
		u8[] texture_data = m3d.texture[i].d[0 .. w * h * ch];
		assert(Transfer(&rd.vk, &model.textures[i], texture_data, w, h), "LoadModel failure: Texture transfer error");

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

	u32[] mat_lookup = AllocArray!(u32)(&rd.temp_arena, m3d.nummaterial);
	model.materials = AllocArray!(Buffer)(&rd.arena, m3d.nummaterial);
	foreach(i; 0 .. m3d.nummaterial)
	{
		Material mat;

		foreach(j; 0 .. m3d.material[i].numprop)
		{
			switch (m3d.material[i].prop[j].type)
			{
			case m3dp_Ka: ConvertColor(&mat.ambient, m3d.material[i].prop[j].value.color); break;
			case m3dp_Ks: ConvertColor(&mat.specular, m3d.material[i].prop[j].value.color); break;
			case m3dp_Ns: mat.shininess = m3d.material[i].prop[j].value.fnum; break;
			case m3dp_map_Kd:
				{
					mat.albedo_texture = tex_lookup[m3d.material[i].prop[j].value.textureid];
					mat.albedo_has_texture = true;
				} break;
			case m3dp_map_Ka:
				{
					mat.ambient_texture = tex_lookup[m3d.material[i].prop[j].value.textureid];
					mat.ambient_has_texture = true;
				} break;
			case m3dp_map_Ks:
				{
					mat.specular_texture = tex_lookup[m3d.material[i].prop[j].value.textureid];
					mat.specular_has_texture = true;
				} break;
			default: break;
			}
		}

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

		mat_lookup[i] = Pop(&rd.vk, 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)(&rd.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)(&rd.temp_arena, m3d.numface * 3);
	Vertex[] vertices = AllocArray!(Vertex)(&rd.temp_arena, m3d.numface * 3);

	foreach(i; 0 .. m3d.numface)
	{
		u32 vi = (i * 3);

		CopyVertex(&vertices[vi+0].pos, &m3d.vertex[m3d.face[i].vertex[0]]);
		CopyVertex(&vertices[vi+1].pos, &m3d.vertex[m3d.face[i].vertex[1]]);
		CopyVertex(&vertices[vi+2].pos, &m3d.vertex[m3d.face[i].vertex[2]]);

		CopyVertex(&vertices[vi+0].n, &m3d.vertex[m3d.face[i].vertex[0]]);
		CopyVertex(&vertices[vi+1].n, &m3d.vertex[m3d.face[i].vertex[1]]);
		CopyVertex(&vertices[vi+2].n, &m3d.vertex[m3d.face[i].vertex[2]]);

		vertices[vi+0].col = m3d.vertex[m3d.face[i].vertex[0]].color;
		vertices[vi+1].col = m3d.vertex[m3d.face[i].vertex[1]].color;
		vertices[vi+2].col = m3d.vertex[m3d.face[i].vertex[2]].color;

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

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

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

		indices[vi+0] = vi+0;
		indices[vi+1] = vi+1;
		indices[vi+2] = vi+2;
	}

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

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

	Reset(&rd.temp_arena);

	return model;
}

void
Destroy(Renderer* rd)
{
	Destroy(&rd.vk, rd.triangle_pipeline);
	Destroy(&rd.vk, rd.compute_pipeline);
	Destroy(&rd.vk);
}