#include "model.hpp"
#include "renderer.hpp"

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

#include <string.h>

struct VP_Cbuffer {
	m4f view_projection;
};

struct TS_Cbuffer {
	m4f inv_view;
	m4f inv_proj;
	m4f prev_vp;
};

struct TS_Caster_Config {
	int index;
	char pad[60];
};

struct Global_Cbuffer {
	v3f camera_pos;
	int light_count;
	int frame;
};

void init_drawlist(
	Drawlist* l,
	Device* d,
	Arena* a,
	int cap
) {
	l->count = 0;
	l->cap = cap;
	l->camera = 0;
	l->vp.init(
		d,
		"Drawlist cbuffer",
		sizeof(VP_Cbuffer),
		Buffer_Flags::constant_buffer
	);
	l->models = (Model_Instance**)arena_alloc(
		a,
		cap * sizeof *l->models
	);
}

void Renderer::make_ts_sampler(Device* d) {
	Sampler_State s{};
	s.min = Filter_Mode::point;
	s.mag = Filter_Mode::point;
	s.mip = Filter_Mode::point;
	s.address_u = Address_Mode::border;
	s.address_v = Address_Mode::border;
	s.border[0] = 0.0f;
	ts_sampler = d->create_sampler("temporal shadow sampler", s);
}

void Renderer::init(
	Arena* arena,
	Device* d,
	Asset_Arena& assets
) {
	int i;
	auto id = [&](int did, int cap) {
		init_drawlist(&drawlists[did], d, arena, cap);
	};
	id(FORWARD, 512);
	for (i = SHADOW_MAP_START; i < SHADOW_MAP_END; i++)
		id(i, 512);
	camera_count = 1;
	cameras.init();
	globals.init(
		d,
		"Renderer globals",
		sizeof(Global_Cbuffer),
		Buffer_Flags::constant_buffer
	);
	ts_config.init(
		d,
		"Temporal shadow cbuffer",
		sizeof(TS_Cbuffer),
		Buffer_Flags::constant_buffer
	);
	ts_config2.init(
		d,
		"Temporal shadow casters",
		sizeof(TS_Caster_Config) * Lighting::max_shadows,
		Buffer_Flags::constant_buffer
	);
	quad.init(d);
	ts_shader = (Shader*)assets.load("ts.csh");
	ts_shadowmap_binding = ts_shader->descriptor_binding("shadowmaps");
	ts_depthmap_binding = ts_shader->descriptor_binding("depthmap");
	ts_prev_binding = ts_shader->descriptor_binding("previous");
	ts_vert_binding = ts_shader->binding_index("verts");
	ts_config_binding = ts_shader->descriptor_binding("config");
	ts_caster_config_binding = ts_shader->descriptor_binding("caster_config");
	ts_casters_binding = ts_shader->descriptor_binding("casters");
	ts_globals_binding = ts_shader->descriptor_binding("globals");
	assert(ts_shadowmap_binding >= 0);
	assert(ts_depthmap_binding >= 0);
	assert(ts_prev_binding >= 0);
	assert(ts_vert_binding >= 0);
	assert(ts_config_binding >= 0);
	assert(ts_caster_config_binding >= 0);
	assert(ts_casters_binding >= 0);
	assert(ts_globals_binding >= 0);
	make_ts_sampler(d);
	frame = 0;
}

void Renderer::destroy(Device* d) {
	int i;
	for (i = 0; i < drawlist_count; i++)
		drawlists[i].vp.destroy(d);
	globals.destroy(d);
	ts_config.destroy(d);
	ts_config2.destroy(d);
	quad.destroy(d);
	d->destroy_sampler(ts_sampler);
}

void Renderer::set_camera(Camera_Id cam, int drawlist) {
	drawlists[drawlist].camera = cam;
}

void Drawlist::render(
	const Renderer& r,
	Device* dev,
	Arena* a,
	const Lighting* l,
	Render_Pass& pass,
	void (*overrider)(Pipeline_Builder&)
) {
	int i, c = count;
	const Camera& cam = r.get_camera(camera);
	Model_Resources res{};
	VP_Cbuffer* vpc = (VP_Cbuffer*)vp.map(dev);
	vpc->view_projection = cam.get_proj() * cam.get_view();
	vp.unmap(dev);
	vp.update(dev->get_ctx());
	res.shadows = l->ss_shadows[r.frame & 1];
	res.sampler = r.clamped_linear;
	res.env_cubemap = r.env_cubemap;
	res.vp = vp.gpuonly;
	res.globals = r.globals.gpuonly;
	res.overrider = overrider;
	for (i = 0; i < c; i++) {
		models[i]->render(dev, a, pass, l, res);
	}
}

void Renderer::update_globals(
	const Lighting* l,
	Device* d,
	Context& ctx
) {
	void* ptr = globals.map(d);
	Global_Cbuffer* cb = (Global_Cbuffer*)ptr;
	v3f& cp = get_camera(drawlists[FORWARD].camera).position;
	cb->camera_pos = cp;
	cb->frame = frame;
	cb->light_count = l->light_count;
	globals.unmap(d);
	globals.update(ctx);
}

void Renderer::temporal_shadows(
	Device* dev,
	Context& ctx,
	const Lighting* l,
	Pipeline_Builder& pb
) {
	int i, c = l->caster_count;
	Camera& cam = get_camera(drawlists[FORWARD].camera);
	TS_Cbuffer* cbuf = (TS_Cbuffer*)ts_config.map(dev);
	TS_Caster_Config* casters = (TS_Caster_Config*)ts_config.map(dev);
	for (i = 0; i < c; i++) {
		casters[i].index = i;
	}
	cbuf->inv_view = cam.get_view().inverse();
	cbuf->inv_proj = cam.get_proj().inverse();
	ts_config.unmap(dev);
	ts_config2.unmap(dev);
	ctx.debug_push("temporal shadows");
	ts_config.update(ctx);
	ts_config2.update(ctx);
	for (i = 0; i < c; i++) {
		int ind = frame & 1;
		auto& pass = pb
			.begin_rp()
			.rp_target(l->ss_shadow_slices[ind][i], Clear_Mode::discard)
			.build_rp();
		auto& pip = pb
			.begin()
			.shader(ts_shader->id)
			.vertex_format(ts_shader->vf)
			.texture(
				ts_shadowmap_binding,
				l->shadows,
				l->shadow_sampler)
			.texture(
				ts_prev_binding,
				l->ss_shadow_slices[!ind][i],
				ts_sampler)
			.texture(
				ts_depthmap_binding,
				dev->get_depth_target(),
				ts_sampler)
			.cbuffer(ts_config_binding, ts_config.gpuonly)
			.cbuffer(
				ts_caster_config_binding,
				ts_config2.gpuonly,
				i * sizeof(TS_Caster_Config),
				sizeof(TS_Caster_Config))
			.cbuffer(ts_globals_binding, globals.gpuonly)
			.sbuffer(ts_casters_binding, l->casters.gpuonly)
			.build();
		quad.render(
			ctx,
			pip,
			pass,
			ts_vert_binding
		);
	}
	ctx.debug_pop();
}

void Renderer::render(
	Device* dev,
	Arena* a, Texture_Id hdr_target,
	const Lighting* l
) {
	int i, j;
	Pipeline_Builder pb(a, dev);
	Context& ctx = dev->get_ctx();

	update_globals(l, dev, ctx);

	Render_Pass& depth_prepass = pb
		.begin_rp()
		.rp_depth_target(dev->get_depth_target(), 1.0f)
		.build_rp();
	Render_Pass& forward_pass = pb
		.begin_rp()
		.rp_target(hdr_target, Clear_Mode::restore)
		.rp_depth_target(dev->get_depth_target(), Clear_Mode::restore)
		.build_rp();
	Render_Pass* shadow_passes[Lighting::max_shadows];
	for (i = 0; i < Lighting::max_shadows; i++) {
		shadow_passes[i] = &pb
			.begin_rp()
			.rp_depth_target(l->shadow_slices[i], 1.0f)
			.build_rp();
	}

	ctx.debug_push("depth prepass");
	drawlists[FORWARD].render(*this, dev, a, l, depth_prepass, 0);
	ctx.debug_pop();

	ctx.debug_push("shadow maps");
	for (
		i = SHADOW_MAP_START, j = 0;
		i < SHADOW_MAP_END && j < l->caster_count;
		i++, j++
	) {
		auto o = [](Pipeline_Builder& pb) {
			pb.cull(Cull_Mode::front);
		};
		setcam(i, l->cameras[j]);
		ctx.debug_push("shadow map");
		drawlists[i].render(*this, dev, a, l, *shadow_passes[j], o);
		ctx.debug_pop();
	}
	ctx.debug_pop();

	temporal_shadows(dev, ctx, l, pb);

	ctx.debug_push("forward");
	drawlists[FORWARD].render(*this, dev, a, l, forward_pass, 0);
	ctx.debug_pop();
	frame++;
}

void Renderer::add_model(int drawlist, Model_Instance* m) {
	Drawlist& d = drawlists[drawlist];
	assert(d.count < d.cap);
	if (d.count >= d.cap) {
		print_war("Drawlist is full.\n");
		return;
	}
	d.models[d.count++] = m;
}

void Renderer::rem_model(int drawlist, Model_Instance* m) {
	int i;
	Drawlist& d = drawlists[drawlist];
	for (i = d.count - 1; i >= 0; i--)
		if (d.models[i] == m)
			d.models[i] = d.models[--d.count];
	assert(0);
}

void Renderer::default_model(Model_Instance* m) {
	int i;
	for (i = 0; i < drawlist_count; i++)
		add_model(i, m);
}

Camera_Id Renderer::create_camera() {
	Camera_Id id(camera_count++);
	Camera cam;
	cameras.set(id, cam);
	return id;
}

Camera& Renderer::get_camera(Camera_Id id) {
	assert(id.index);
	return cameras[id];
}

const Camera& Renderer::get_camera(Camera_Id id) const {
	assert(id.index);
	return cameras[id];
}

void Renderer::destroy_camera(Camera_Id cam) {
	cameras.remove(cam);
}

void Renderer::setcam(int did, Camera_Id cam) {
	assert(cam.index);
	drawlists[did].camera = cam;
}

void Fullscreen_Quad::init(Device* d) {
	float verts[] = {
		-1.0f, -1.0f, 0.0f, 0.0f,
		-1.0f,  3.0f, 0.0f, 2.0f,
		 3.0f, -1.0f, 2.0f, 0.0f
	};
	Buffer_Id stage;
	void* mem;
	stage = d->create_buffer(
		"sky vb stage",
		sizeof verts,
		Buffer_Flags::cpu_readwrite |
		Buffer_Flags::copy_src
	);
	mem = d->map_buffer(stage, 0, sizeof verts);
	memcpy(mem, verts, sizeof verts);
	d->unmap_buffer(stage);
	vb = d->create_buffer(
		"fullscreen quad",
		sizeof verts,
		Buffer_Flags::copy_dst |
		Buffer_Flags::vertex_buffer
	);
	Context& ctx = d->acquire();
	ctx.copy(vb, stage);
	d->submit(ctx);
	d->destroy_bufferi(stage);
}

void Fullscreen_Quad::destroy(Device* d) {
	d->destroy_buffer(vb);
}

void Fullscreen_Quad::render(
	Context& ctx,
	Pipeline& pip,
	Render_Pass& pass,
	int bind
) {
	Vertex_Buffer_Binding vbb[] = {{
		.id = vb,
		.offset = 0,
		.target = bind
	}, {}};
	Draw draw{};
	draw.verts = vbb;
	draw.vertex_count = 3;
	draw.instance_count = 1;
	ctx.submit(draw, pip, pass);
}