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 u = util;
import p = platform;

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

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,
}

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 triangle_pipeline;
	Pipeline compute_pipeline;
	Pipeline ui_pipeline;
}

struct GlobalUniforms
{
	Vec2 res;
}

struct ShaderUniforms
{
	f32 placeholder;
}

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

struct Vertex
{
	Vec3 pos;
	Vec3 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;
	u32[]				m_indices;

	ImageView[] textures;
	u32[]			 	t_indices;
}

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 },
		],
	};

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

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

		Bind(rd, &rd.triangle_pipeline);

		Draw(rd, 3, 1);

		success = FinishFrame(rd);
	}

	return success;
}

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

	return model;
}

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

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

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

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

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

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

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

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

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

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

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

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

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


