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

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

#include <string.h>

struct MVP_Cbuffer {
	m4f model;
	m4f view_projection;
};

struct Mat_Cbuffer {
	float albedo[3];
	float metalness;
	float roughness;
	float ao;
	char pad2[8];
	float camera_pos[3];
	char pad[20];
};

void Material::use(
	Pipeline_Builder& pb,
	Sampler_Id sampler,
	Shader& shader
) {
	auto bind = [&shader, &pb, sampler](
		const char* name,
		Texture_Id t
	) {
		int loc = shader.descriptor_binding(name);
		if (loc >= 0) {
			pb.texture(loc, t, sampler);
		}
	};
	bind("albedo", tex.albedo);
	bind("ao",     tex.ao);
	bind("metal",  tex.metal);
	bind("rough",  tex.rough);
	bind("normal", tex.normal);
}

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];
		char depth_shader_name[28];
		char material_name[28];
		Shader* shader, * depth_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);
		pack_read(f, depth_shader_name, sizeof depth_shader_name);
		pack_read(f, material_name, sizeof material_name);
		mesh.offset = icoff;
		shader = (Shader*)shaders->load(shader_name);
		depth_shader = (Shader*)shaders->load(depth_shader_name);
		mesh.material = (Material*)shaders->load(material_name);
		assert(shader != 0);
		assert(depth_shader != 0);
		assert(mesh.material != 0);
		mesh.shader = shader->id;
		mesh.depth_shader = depth_shader->id;
		mesh.mvp_binding = shader->descriptor_binding("c_mvp");
		mesh.mvp_binding_depth = depth_shader->descriptor_binding("c_mvp");
		mesh.mat_binding = shader->descriptor_binding("material");
		mesh.env_cube_binding = shader->descriptor_binding("env_cube");
		mesh.mesh_binding = shader->binding_index("mesh");
		assert(mesh.mvp_binding >= 0);
		assert(mesh.mat_binding >= 0);
		assert(mesh.mesh_binding >= 0);
		assert(mesh.env_cube_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, &mesh.bound.min, 12);
		pack_read(f, &mesh.bound.max, 12);
		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
	);
	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);
	r->update_transforms();
	return r;
}

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

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 Material_Loader::init(
	Asset_Arena* texture_arena,
	Texture_Id dt
) {
	default_tex = dt;
	textures = texture_arena;
}

Asset* Material_Loader::load(
		Arena* a,
		Arena* s,
		const char* filename,
		Pack_File* f
) {
	Material_File mf;
	Material* m = (Material*)arena_alloc(a, sizeof *m);
	auto read_name = [&](int len) {
		char* r = (char*)arena_alloc(s, len + 1);
		r[pack_read(f, r, len)] = 0;
		return r;
	};
	auto read_tex = [&](int len) {
		if (!len) return default_tex;
		const char* name = read_name(len);
		auto t = (Texture*)textures->load(name);
		if (!t)
			return default_tex;
		return t->id;
	};
	(void)filename;
	pack_read(f, &mf, sizeof mf);
	if (
		mf.magic[0] != 'M' ||
		mf.magic[1] != 'T' ||
		mf.magic[2] != 'R' ||
		mf.magic[3] != 'L'
	) return 0;
	m->metalness = mf.metalness;
	m->roughness = mf.roughness;
	m->ao = mf.ao;
	m->albedo = v3f(
		(float)((mf.albedo >> 16) & 0xff) / 255.0f,
		(float)((mf.albedo >> 8)  & 0xff) / 255.0f,
		(float)((mf.albedo)       & 0xff) / 255.0f
	);
	m->tex.albedo = read_tex(mf.albedo_tex_len);
	m->tex.ao = read_tex(mf.ao_tex_len);
	m->tex.metal = read_tex(mf.metal_tex_len);
	m->tex.rough = read_tex(mf.rough_tex_len);
	m->tex.normal = read_tex(mf.normal_tex_len);
	return m;
}

void Material_Loader::unload(Asset* a) {
	(void)a;
}

void Camera::init(float vfov, const v3f& f, const v3f& p) {
	fov = vfov;
	forward = f;
	position = p;
	near = 0.1f;
	far = 1000.0f;
	asp = 1.0f;
}

m4f Camera::get_view() const {
	v3f up(0.0f, 1.0f, 0.0f);
	return m4f::lookat(position, position + forward, up);
}

m4f Camera::get_proj() const {
	return m4f::pers(fov, asp, near, far);
}

void Model_Instance::init(Device* dev, Model* model) {
	m = model;
	mvp = dev->create_buffer(
		"Model instance MVP",
		sizeof(MVP_Cbuffer) * m->mesh_count,
		Buffer_Flags::constant_buffer |
		Buffer_Flags::cpu_readwrite
	);
	mat = dev->create_buffer(
		"Model instance material",
		sizeof(Mat_Cbuffer) * m->mesh_count,
		Buffer_Flags::constant_buffer |
		Buffer_Flags::cpu_readwrite
	);
}

void Model_Instance::destroy(Device* dev) {
	dev->destroy_buffer(mat);
	dev->destroy_buffer(mvp);
}

void Model_Instance::update_cbuffers(
	Device* dev,
	const Camera& cam
) {
	int i, c = m->mesh_count;
	Mesh* meshes = m->get_meshes();
	MVP_Cbuffer* mvps = (MVP_Cbuffer*)dev->map_buffer(
		mvp,
		0,
		c * sizeof *mvps
	);
	Mat_Cbuffer* mats = (Mat_Cbuffer*)dev->map_buffer(
		mat,
		0,
		c * sizeof *mats
	);
	m4f view_projection = cam.get_proj() * cam.get_view();
	for (i = 0; i < c; i++) {
		Mat_Cbuffer& mat = mats[i];
		Material& sm = *meshes[i].material;
		mvps[i].view_projection = view_projection;
		mvps[i].model = transform * meshes[i].world;
		mat.albedo[0] = sm.albedo.x;
		mat.albedo[1] = sm.albedo.y;
		mat.albedo[2] = sm.albedo.z;
		mat.metalness = sm.metalness;
		mat.roughness = sm.roughness;
		mat.ao = sm.ao;
		mat.camera_pos[0] = cam.position.x;
		mat.camera_pos[1] = cam.position.y;
		mat.camera_pos[2] = cam.position.z;
	}
	dev->unmap_buffer(mvp);
	dev->unmap_buffer(mat);
}

void Model_Instance::render(
	Device* dev,
	Arena* a,
	Render_Pass& pass,
	Texture_Id env_cubemap,
	Sampler_Id sampler
) {
	int i, c = m->mesh_count;
	Mesh* meshes = m->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 = m->vbo,
			.offset = 0,
			.target = mesh.mesh_binding
		}, {}};
		Index_Buffer_Binding ibb = {
			.id = m->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);
			pb.shader(mesh.depth_shader);
			pb.cbuffer(
				mesh.mvp_binding_depth,
				mvp,
				i * sizeof(MVP_Cbuffer),
				sizeof(MVP_Cbuffer)
			);
		} else {
			pb.depth(true, false, Depth_Mode::equal);
			pb.shader(mesh.shader);
			mesh.material->use(pb, sampler, dev->get_shader(mesh.shader));
			pb.cbuffer(
				mesh.mat_binding,
				mat,
				i * sizeof(Mat_Cbuffer),
				sizeof(Mat_Cbuffer)
			);
			pb.cbuffer(
				mesh.mvp_binding,
				mvp,
				i * sizeof(MVP_Cbuffer),
				sizeof(MVP_Cbuffer)
			);
			pb.texture(
				mesh.env_cube_binding,
				env_cubemap,
				sampler
			);
		}
		pb.cull(Cull_Mode::back);
		pb.vertex_format(shader.vf);
		Pipeline& pip = pb.build();
		ctx.submit(draw, pip, pass);
	}
}

void Model_Scene::init(
	Arena* arena,
	int max_instances,
	Sampler_Id s
) {
	instances = (Model_Instance*)arena_alloc(arena, max_instances);
	count = 0;
	max = max_instances;
	sampler = s;
}

Model_Instance* Model_Scene::instantiate(
	Device* dev,
	Model* model
) {
	Model_Instance* instance = &instances[count++];
	assert(count <= max);
	instance->init(dev, model);
	return instance;
}

void Model_Scene::uninstantiate(
	Device* dev,
	Model_Instance* model
) {
	int idx = model - instances;
	int last = count - 1;
	model->destroy(dev);
	instances[idx] = instances[last];
	count = last;
}

void Model_Scene::update(const Camera& cam, Device* dev) {
	int i;
	Model_Instance* instance = instances;
	for (i = 0; i < count; i++, instance++)
		instance->update_cbuffers(dev, cam);
}

void Model_Scene::render(
	Device* dev,
	Arena* a,
	Render_Pass& pass,
	Texture_Id env_cubemap,
	Sampler_Id sampler
) {
	int i;
	Model_Instance* instance = instances;
	for (i = 0; i < count; i++, instance++)
		instance->render(dev, a, pass, env_cubemap, sampler);
}

void Model_Scene::destroy(Device* dev) {
	int i;
	Model_Instance* instance = instances;
	for (i = 0; i < count; i++, instance++)
		instance->destroy(dev);
}