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

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

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,
	RGBA_SRGB = VK_FORMAT_R8G8B8A8_SRGB,
}

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 Specialization
{
	u64					size;
	SpecEntry[] entries;
	void*				data;
}

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

struct CompPipelineInfo
{
	string 			shader;
	Specialization spec;
	VkPipelineLayout *layout;
}

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;

	PushConst push_const;

	Vec3 camera_pos = Vec3(0.0);

	Model yoder;
}

struct Camera
{
	Vec3 pos = Vec3(0.0);
	Vec3 target = Vec3(0.0);
}

struct GlobalUniforms
{
	Mat4 world_matrix;
	Vec2 res;
}

struct ShaderUniforms
{
	f32 placeholder;
}

extern(C) struct Material
{
	Vec4 ambient;
	Vec4 diffuse;
	Vec4 specular;
	u32 albedo_texture;
	u32 ambient_texture;
	u32 specular_texture;
	b32 albedo_has_texture;
	b32 ambient_has_texture;
	b32 specular_has_texture;
	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.spv",
		frag_shader: "shaders/triangle.frag.spv",
	};

	GfxPipelineInfo ui_info = {
		vertex_shader: "shaders/gui.vert.spv",
		frag_shader: "shaders/gui.frag.spv",
		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.spv",
		frag_shader: "shaders/pbr.frag.spv",
		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 },
		],
	};

	CompPipelineInfo gradient_info = {
		shader: "shaders/gradient.comp.spv",
	};

	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, &gradient_info);

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

	ReadModel(&rd, "models/Tree01.m3d");

	return rd;
}

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

	f32 radius = 10.0;
	f32 cam_x = sin(RDTSC()) * radius * 0.000000000001;
	f32 cam_z = cos(RDTSC()) * radius * 0.000000000001;

	rd.globals.world_matrix = Mat4(1.0);

	BeginFrame(rd);

	Bind(rd, &rd.compute_pipeline);

	SetUniform(rd, &rd.globals);

	//DrawRect(rd, 150.0, 300.0, 500.0, 700.0, Vec4(0.0, 0.0, 1.0, 1.0));

	PrepComputeDrawImage(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);

	FinishFrame(rd);
}

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

	foreach(i, part; model.parts)
	{
		rd.push_const.mat_id = part.mat;
		PushConstants(&rd.vk, &rd.push_const);
		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;
	}
}

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

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

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

pragma(inline): void
PrepComputeDrawImage(Renderer* rd)
{
	PrepComputeDrawImage(&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, CompPipelineInfo* comp_info)
{
	return CreateComputePipeline(&rd.vk, comp_info);
}

void
ReadModel(Renderer* rd, string name)
{
	AssetInfo info = GetAssetInfo(name);

	u8[] data = LoadAssetData(&rd.temp_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)];
		Logf("texture name: %r", tex_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)];
		Logf("material name: %r", mat_name);
	}
}

Model
LoadModel(Renderer* rd, string name)
{
	AssetInfo info = GetAssetInfo(name);

	u8[] data = LoadAssetData(&rd.temp_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)(&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;
		u8[] tex_data = m3d.texture[i].d[0 .. w * h * ch];

		Logf("w: %s h: %s ch: %s", w, h, ch);

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

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

	Material[] mats = AllocArray!(Material)(&rd.temp_arena, m3d.nummaterial);
	u32[] mat_lookup = AllocArray!(u32)(&rd.temp_arena, m3d.nummaterial);
	model.materials = AllocArray!(Buffer)(&rd.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_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_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;
			default: break;
			}
		}

		CreateBuffer(&rd.vk, &model.materials[i], BT.Uniform, Material.sizeof, false);
		assert(Transfer(&rd.vk, &model.materials[i], &mats[i]), "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;
	}

	CreateBuffer(&rd.vk, &model.vertex_buffer, BT.Vertex, vertices.length * Vertex.sizeof, false);
	CreateBuffer(&rd.vk, &model.index_buffer, 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");

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

	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);
}


