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

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

#include <string.h>

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

Asset* Model_Loader::load(Arena* a, Arena* s, 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(
		vbo_size,
		Buffer_Flags::copy_src |
		Buffer_Flags::cpu_readwrite
	);
	stage_indices = dev->create_buffer(
		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;
		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;
		}
		mesh.offset = icoff;
		mesh.parent = 0;
		mesh.shader = ((Shader*)shaders->load("surface.csh"))->id;
		pack_read(f, &vertex_size, 4);
		pack_read(f, &mesh.count, 4);
		pack_read(f, &vertex_count, 4);
		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(
		vbo_size,
		Buffer_Flags::copy_dst |
		Buffer_Flags::vertex_buffer
	);
	r->ibo = dev->create_buffer(
		ibo_size,
		Buffer_Flags::copy_dst |
		Buffer_Flags::index_buffer
	);
	ctx = &dev->acquire();
	ctx->copy(r->vbo, stage_verts);
	ctx->copy(r->ibo, stage_indices);
	dev->submit(*ctx);
	dev->destroy_buffer(stage_verts);
	dev->destroy_buffer(stage_indices);
	return r;
}

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

void Model::render(
	Device* dev,
	Arena* a,
	Render_Pass& pass,
	Buffer_Id config
) {
	int i, c = mesh_count;
	Mesh* meshes = get_meshes();
	Context& ctx = dev->get_ctx();
	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 = shader.binding_index("mesh")
		}, {}};
		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();
		pb.depth(true, true, Depth_Mode::less);
		pb.shader(mesh.shader);
		pb.cbuffer(
			shader.descriptor_binding("config_buffer"),
			config
		);
		pb.vertex_format(shader.vf);
		Pipeline& pip = pb.build();
		ctx.submit(draw, pip, pass);
	}
}