#include "lighting.hpp"
#include "model.hpp"
#include "renderer.hpp"
#include "world.hpp"

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

#include <random>

/* needs to match surface shader */
struct GPU_Light {
	v3f dir;
	float brightness;
	v3f colour;
	int caster_id;
};

struct GPU_Caster {
	m4f projection;
	m4f jittered;
};

void Lighting::init(Device* dev, int w, int h) {
	int i;
	Sampler_State ss{};
	gpu_lights.init(
		dev,
		"Light buffer",
		max_lights * sizeof(GPU_Light),
		Buffer_Flags::storage_buffer
	);
	gpu_casters.init(
		dev,
		"Caster buffer",
		max_shadows * sizeof(GPU_Caster),
		Buffer_Flags::storage_buffer
	);
	shadows = dev->create_texture(
		"Shadowmap Array",
		texture_format_d16,
		Texture_Flags::sampleable | Texture_Flags::depth_stencil_target,
		shadow_res,
		shadow_res,
		1,
		1,
		shadowmap_count,
		0
	);
	zero(ss_shadows, sizeof ss_shadows);
	zero(ss_shadow_slices, sizeof ss_shadow_slices);
	zero(occlusion, sizeof occlusion);
	recreate(dev, w, h);
	for (i = 0; i < shadowmap_count; i++) {
		cameras[i] = 0;
		shadow_slices[i] = dev->alias_texture(
			shadows,
			"Shadowmap Slice",
			texture_format_d16,
			Texture_Flags::depth_stencil_target,
			shadow_res,
			shadow_res,
			1,
			1,
			1,
			0,
			i
		);
	}
	ss.min = Filter_Mode::linear;
	ss.mag = Filter_Mode::linear;
	ss.mip = Filter_Mode::linear;
	ss.address_u = Address_Mode::border;
	ss.address_v = Address_Mode::border;
	ss.border[0] = 1.0f;
	ss.border[1] = 1.0f;
	ss.border[2] = 1.0f;
	ss.border[3] = 1.0f;
	ss.compare = Depth_Mode::less;
	shadow_sampler = dev->create_sampler("shadow sampler", ss);
}

void Lighting::destroy(Device* dev, Renderer& r) {
	int i;
	gpu_lights.destroy(dev);
	gpu_casters.destroy(dev);
	dev->destroy_sampler(shadow_sampler);
	for (i = 0; i < shadowmap_count; i++) {
		dev->destroy_texture(shadow_slices[i]);
		r.destroy_camera(cameras[i]);
	}
	dev->destroy_texture(shadows);
	destroy_ss(dev);
}

void Lighting::destroy_ss(Device* dev) {
	int i, j;
	for (i = 0; i < 2; i++) {
		for (j = 0; j < max_shadows; j++)
			if (ss_shadow_slices[i][j])
				dev->destroy_texture(ss_shadow_slices[i][j]);
		if (ss_shadows[i])
			dev->destroy_texture(ss_shadows[i]);
		if (occlusion[i])
			dev->destroy_texture(occlusion[i]);
	}
}

void Lighting::recreate(Device* dev, int w, int h) {
	int i, j;
	destroy_ss(dev);
	for (i = 0; i < 2; i++) {
		ss_shadows[i] = dev->create_texture(
			"Shadow accumulation buffer",
			texture_format_r16f,
			Texture_Flags::sampleable | Texture_Flags::colour_target,
			w,
			h,
			1,
			1,
			max_shadows,
			0
		);
		for (j = 0; j < max_shadows; j++)
			ss_shadow_slices[i][j] = dev->alias_texture(
				ss_shadows[i],
				"Shadow accumulation buffer slice",
				texture_format_r16f,
				Texture_Flags::colour_target,
				w,
				h,
				1,
				1,
				1,
				0,
				j
			);
		occlusion[i] = dev->create_texture(
			"Occlusion buffer",
			texture_format_r32f,
			Texture_Flags::sampleable | Texture_Flags::colour_target,
			w,
			h,
			1,
			1,
			1,
			0
		);
	}
}

Camera_Id Lighting::gm_cam(Renderer& r) {
	int id = cam_count++;
	Camera_Id cid = cameras[id];
	if (!cid) {
		cid = r.create_camera();
		cameras[id] = cid;
	}
	return cid;
}

void Lighting::write_bufs(
	void* lptr,
	void* cptr,
	World& w,
	Renderer& r,
	Model_Scene& s
) {
	GPU_Light* ldst  = (GPU_Light*)lptr;
	GPU_Caster* cdst = (GPU_Caster*)cptr;
	int count = 0, ccount = 0;
	cam_count = 0;
	std::random_device dev;
	std::mt19937 rng(dev());
	std::uniform_real_distribution<float> dist6(-0.01f, 0.01f);
	for (auto v : w.view<Sun_Light>()) {
		GPU_Light gl;
		Sun_Light& l = v.get<Sun_Light>();
		if (count >= max_lights) {
			print_war("Over light limit.\n");
			return;
		}
		gl.brightness = l.brightness;
		gl.colour = l.colour;
		gl.dir = l.dir;
		if (l.caster && ccount < max_shadows) {
			Camera* cam[2];
			int cid = ccount++;
			Caster& caster = casters[cid];
			GPU_Caster& gc = cdst[cid];
			v3f jitter(
				dist6(rng),
				dist6(rng),
				dist6(rng)
			);
			caster = Caster { gm_cam(r), gm_cam(r) };
			cam[0] = &r.get_camera(caster.reality);
			cam[1] = &r.get_camera(caster.jittered);
			cam[0]->init_shadow(l.dir,          s.bound.min, s.bound.max);
			cam[1]->init_shadow(l.dir + jitter, s.bound.min, s.bound.max);
			gc.projection = cam[0]->get_proj() * cam[0]->get_view();
			gc.jittered   = cam[1]->get_proj() * cam[1]->get_view();
			gl.caster_id = cid;
		} else
			gl.caster_id = -1;
		ldst[count++] = gl;
	}
	light_count = count;
	caster_count = ccount;
}

void Lighting::update(
	Device* dev,
	Context& ctx,
	World& w,
	Renderer& r,
	Model_Scene& s
) {
	light_count = 0;
	write_bufs(gpu_lights.map(dev), gpu_casters.map(dev), w, r, s);
	gpu_casters.unmap(dev);
	gpu_lights.unmap(dev);
	gpu_lights.update(ctx);
	gpu_casters.update(ctx);
}