#include "maths.hpp"
#include "model.hpp"

extern "C" {
#include "memory.h"
#include "pack.h"
#include "plat.h"
}

#include <string.h>

struct MVP_Cbuffer {
	m4f model;
	m4f view_projection;
};

void Model_Loader::init(Device* device, Asset_Arena* shader_arena) {
	dev = device;
	shaders = shader_arena;
}

Asset* Model_Loader::load(
	Arena* a,
	Arena* s,
	const char* filename,
	Pack_File* f
) {
	char magic[4];
	int vbo_size, ibo_size, node_count, i, coff, icoff, vcoff;
	uint8_t* verts;
	uint16_t* indices;
	Model* r;
	Mesh* meshes;
	Context* ctx;
	Buffer_Id stage_verts, stage_indices;
	(void)s;
	pack_read(f, magic, 4);
	if (
		magic[0] != 'M' ||
		magic[1] != 'O' ||
		magic[2] != 'D' ||
		magic[3] != 'L'
	) {
		print_err("Invalid model.\n");
		return 0;
	}
	pack_read(f, &vbo_size, 4);
	pack_read(f, &ibo_size, 4);
	pack_read(f, &node_count, 4);
	r = (Model*)arena_alloc(a, sizeof *r + node_count * sizeof(Mesh));
	r->mesh_count = node_count;
	meshes = r->get_meshes();
	stage_verts = dev->create_buffer(
		"model vertex buffer stage",
		vbo_size,
		Buffer_Flags::copy_src |
		Buffer_Flags::cpu_readwrite
	);
	stage_indices = dev->create_buffer(
		"model index buffer stage",
		ibo_size,
		Buffer_Flags::copy_src |
		Buffer_Flags::cpu_readwrite
	);
	verts = (uint8_t*)dev->map_buffer(stage_verts, 0, vbo_size);
	indices = (uint16_t*)dev->map_buffer(stage_indices, 0, ibo_size);
	for (i = 0, coff = 0, icoff = 0, vcoff = 0; i < node_count; i++) {
		Mesh& mesh = meshes[i];
		int vertex_size, vertex_count;
		char shader_name[28];
		Shader* shader;
		pack_read(f, magic, 4);
		if (
			magic[0] != 'M' ||
			magic[1] != 'E' ||
			magic[2] != 'S' ||
			magic[3] != 'H'
		) {
			print_err("Invalid mesh.\n");
			return 0;
		}
		pack_read(f, shader_name, sizeof shader_name);
		mesh.offset = icoff;
		shader = (Shader*)shaders->load(shader_name);
		assert(shader != 0);
		mesh.shader = shader->id;
		mesh.mvp_binding = shader->descriptor_binding("c_mvp");
		mesh.mesh_binding = shader->binding_index("mesh");
		assert(mesh.mvp_binding >= 0);
		assert(mesh.mesh_binding >= 0);
		pack_read(f, &vertex_size, 4);
		pack_read(f, &mesh.count, 4);
		pack_read(f, &vertex_count, 4);
		pack_read(f, &mesh.parent, 4);
		pack_read(f, &mesh.local, 64);
		pack_read(f, &verts[coff], vertex_count * vertex_size);
		pack_read(f, &indices[icoff], mesh.count * sizeof *indices);
		mesh.vbo_offset = vcoff;
		icoff += mesh.count;
		coff += vertex_size * vertex_count;
		vcoff += vertex_count;
	}
	dev->unmap_buffer(stage_verts);
	dev->unmap_buffer(stage_indices);
	r->vbo = dev->create_buffer(
		filename,
		vbo_size,
		Buffer_Flags::copy_dst |
		Buffer_Flags::vertex_buffer
	);
	r->ibo = dev->create_buffer(
		filename,
		ibo_size,
		Buffer_Flags::copy_dst |
		Buffer_Flags::index_buffer
	);
	r->mvp = dev->create_buffer(
		filename,
		sizeof(MVP_Cbuffer) * node_count,
		Buffer_Flags::constant_buffer |
		Buffer_Flags::cpu_readwrite
	);
	ctx = &dev->acquire();
	ctx->copy(r->vbo, stage_verts);
	ctx->copy(r->ibo, stage_indices);
	dev->submit(*ctx);
	dev->destroy_bufferi(stage_verts);
	dev->destroy_bufferi(stage_indices);
	return r;
}

void Model_Loader::unload(Asset* a) {
	Model* m = (Model*)a;
	dev->destroy_buffer(m->vbo);
	dev->destroy_buffer(m->ibo);
	dev->destroy_buffer(m->mvp);
}

const m4f& Mesh::get_world(Model& m) {
	if (world_dirty) {
		if (parent)
			world = m.get_meshes()[parent].get_world(m) * local;
		else
			world = local;
		world_dirty = false;
	}
	return world;
}

void Model::update_transforms() {
	int i, c = mesh_count;
	Mesh* meshes = get_meshes();
	for (i = 0; i < c; i++)
		meshes[i].world_dirty = true;
	for (i = 0; i < c; i++)
		meshes[i].get_world(*this);
}

void Model::update_cbuffers(
	Device* dev,
	const m4f& transform,
	const m4f& view_projection
) {
	int i, c = mesh_count;
	Mesh* meshes = get_meshes();
	MVP_Cbuffer* cbuffers = (MVP_Cbuffer*)dev->map_buffer(
		mvp,
		0,
		c * sizeof *cbuffers
	);
	for (i = 0; i < c; i++) {
		cbuffers[i].view_projection = view_projection;
		cbuffers[i].model = transform * meshes[i].world;
	}
	dev->unmap_buffer(mvp);
}

void Model::render(
	Device* dev,
	Arena* a,
	Render_Pass& pass
) {
	int i, c = mesh_count;
	Mesh* meshes = get_meshes();
	Context& ctx = dev->get_ctx();
	bool depth_only = !pass.colour_count;
	for (i = 0; i < c; i++) {
		Mesh& mesh = meshes[i];
		Shader& shader = dev->get_shader(mesh.shader);
		Vertex_Buffer_Binding vbb[] = {{
			.id = vbo,
			.offset = 0,
			.target = mesh.mesh_binding
		}, {}};
		Index_Buffer_Binding ibb = {
			.id = ibo,
			.offset = 0
		};
		Draw draw{};
		Pipeline_Builder pb(a, dev);
		draw.verts = vbb;
		draw.inds = ibb;
		draw.vertex_count = mesh.count;
		draw.instance_count = 1;
		draw.first_vertex = mesh.offset;
		draw.vertex_offset = mesh.vbo_offset;
		pb.begin();
		if (depth_only)
			pb.depth(true, true, Depth_Mode::less);
		else
			pb.depth(true, false, Depth_Mode::equal);
		pb.cull(Cull_Mode::back);
		pb.shader(mesh.shader);
		pb.cbuffer(
			mesh.mvp_binding,
			mvp,
			i * sizeof(MVP_Cbuffer),
			sizeof(MVP_Cbuffer)
		);
		pb.vertex_format(shader.vf);
		Pipeline& pip = pb.build();
		ctx.submit(draw, pip, pass);
	}
}