#include <psemek/gfx/renderer/deferred.hpp>

#include <psemek/gfx/program.hpp>
#include <psemek/gfx/framebuffer.hpp>
#include <psemek/gfx/texture.hpp>
#include <psemek/gfx/mesh.hpp>
#include <psemek/gfx/error.hpp>

#include <psemek/geom/homogeneous.hpp>

#include <psemek/cg/convex_hull_2d/graham.hpp>

namespace psemek::gfx
{

	char const g_buffer_pass_common[] =
R"(#version 330

const uint O_UNIFORM_COLOR  = 1u << 0;
const uint O_TEXTURE_COLOR  = 1u << 1;
const uint O_TRANSPARENT    = 1u << 2;
const uint O_LIT            = 1u << 3;
const uint O_CASTS_SHADOW   = 1u << 4;
const uint O_PRE_TRANSFORM  = 1u << 5;
const uint O_POST_TRANSFORM = 1u << 6;
const uint O_INSTANCED      = 1u << 7;

uniform uint u_flag_mask;
)";

	char const g_buffer_pass_vs[] =
R"(

uniform mat4 u_camera_transform;
uniform mat4x3 u_pre_transform;
uniform mat4x3 u_post_transform;

layout (location = 0) in vec4 in_position;
layout (location = 1) in vec4 in_color;
layout (location = 2) in vec2 in_texcoord;
layout (location = 3) in vec3 in_normal;

layout (location = 4) in mat3x4 in_instance_transform;

out vec3 position;
out vec4 color;
out vec2 texcoord;
out vec3 normal;

void main()
{
	vec4 pos = in_position;
	vec3 n = in_normal;

	if ((u_flag_mask & O_PRE_TRANSFORM) != 0u)
	{
		pos = vec4(u_pre_transform * pos, 1.0);
		n = u_pre_transform * vec4(n, 0.0);
	}

	if ((u_flag_mask & O_INSTANCED) != 0u)
	{
		pos = vec4(transpose(in_instance_transform) * pos, 1.0);
		n = transpose(in_instance_transform) * vec4(n, 0.0);
	}

	if ((u_flag_mask & O_POST_TRANSFORM) != 0u)
	{
		pos = vec4(u_post_transform * pos, 1.0);
		n = u_post_transform * vec4(n, 0.0);
	}

	position = pos.xyz;

	gl_Position = u_camera_transform * pos;

	color = in_color;
	texcoord = in_texcoord;
	normal = n;
}
)";

	char const g_buffer_pass_fs[] =
R"(

uniform vec4 u_color;
uniform sampler2D u_texture;
uniform vec3 u_material;
uniform float u_max_intensity;

in vec3 position;
in vec4 color;
in vec2 texcoord;
in vec3 normal;

layout (location = 0) out vec3 out0;
layout (location = 1) out vec4 out1;
layout (location = 2) out vec3 out2;
layout (location = 3) out vec3 out3;

void main()
{
	vec4 albedo;
	if ((u_flag_mask & O_TEXTURE_COLOR) != 0u)
	{
		vec4 base_color = texture(u_texture, texcoord);
		if ((u_flag_mask & O_UNIFORM_COLOR) != 0u)
			albedo = u_color * base_color;
		else
			albedo = base_color;
	}
	else
	{
		if ((u_flag_mask & O_UNIFORM_COLOR) != 0u)
			albedo = u_color;
		else
			albedo = color;
	}

	out0 = position;
	out1 = vec4(albedo.rgb / u_max_intensity, (u_flag_mask & O_LIT) != 0u ? 1.f : 0.f);
	out2 = normalize(normal) * 0.5 + vec3(0.5);
	out3 = u_material;
}
)";

	char const fullscreen_vs[] =
R"(#version 330

const vec4 vertices[6] = vec4[6](
	vec4(-1.0, -1.0, 0.0, 1.0),
	vec4( 1.0, -1.0, 0.0, 1.0),
	vec4( 1.0,  1.0, 0.0, 1.0),

	vec4(-1.0, -1.0, 0.0, 1.0),
	vec4( 1.0,  1.0, 0.0, 1.0),
	vec4(-1.0,  1.0, 0.0, 1.0)
);

out vec2 texcoord;

void main()
{
	gl_Position = vertices[gl_VertexID];

	texcoord = vertices[gl_VertexID].xy * 0.5 + vec2(0.5);
}
)";

	char const screen_vs[] =
R"(#version 330

layout (location = 0) in vec2 in_position;

out vec2 texcoord;

void main()
{
	gl_Position = vec4(in_position, 0.0, 1.0);

	texcoord = in_position * 0.5 + vec2(0.5);
}
)";

	char const ambient_pass_fs[] =
R"(#version 330

uniform sampler2D u_g0;
uniform sampler2D u_g1;
uniform sampler2D u_g2;
uniform sampler2D u_g3;

uniform vec3 u_ambient;

uniform float u_max_intensity;

in vec2 texcoord;

out vec4 out_color;

void main()
{
	vec4 albedo = texture(u_g1, texcoord);

	vec3 color;

	if (albedo.a < 0.5)
		color = albedo.rgb;
	else
		color = albedo.rgb * u_ambient;

	out_color = vec4(color, 1.0);
}
)";

	char const directional_light_pass_fs[] =
R"(#version 330

uniform sampler2D u_g0;
uniform sampler2D u_g1;
uniform sampler2D u_g2;
uniform sampler2D u_g3;

uniform vec3 u_light_direction;
uniform vec3 u_light_color;

uniform vec3 u_camera_position;

uniform float u_max_intensity;

in vec2 texcoord;

out vec4 out_color;

void main()
{
	vec3 position = texture(u_g0, texcoord).xyz;
	vec4 albedo = texture(u_g1, texcoord);
	vec3 normal = texture(u_g2, texcoord).xyz * 2.0 - vec3(1.0);
	vec3 material = texture(u_g3, texcoord).xyz;

	vec3 view = normalize(u_camera_position - position);

	float d = dot(u_light_direction, normal);

	vec3 refl = 2.0 * normal * d - u_light_direction;

	float l = max(0.0, d) * material.x + pow(max(0.0, dot(view, refl)), material.z) * material.y;

	vec3 color = l * albedo.rgb * u_light_color * albedo.a;

	out_color = vec4(color, 1.0);
}
)";

	char const point_light_pass_fs[] =
R"(#version 330

uniform sampler2D u_g0;
uniform sampler2D u_g1;
uniform sampler2D u_g2;
uniform sampler2D u_g3;

uniform vec3 u_light_position;
uniform vec3 u_light_color;
uniform vec3 u_light_attenuation;

uniform vec3 u_camera_position;

uniform float u_max_intensity;

in vec2 texcoord;

out vec4 out_color;

void main()
{
	vec3 position = texture(u_g0, texcoord).xyz;
	vec4 albedo = texture(u_g1, texcoord);
	vec3 normal = texture(u_g2, texcoord).xyz * 2.0 - vec3(1.0);
	vec3 material = texture(u_g3, texcoord).xyz;

	vec3 view = normalize(u_camera_position - position);

	vec3 light = u_light_position - position;

	float r = length(light);
	light /= r;

	float d = dot(light, normal);

	vec3 refl = 2.0 * normal * d - light;

	float l = max(0.0, d) * material.x + pow(max(0.0, dot(view, refl)), material.z) * material.y;

	vec3 color = l * albedo.rgb * u_light_color * albedo.a / (u_light_attenuation.x + r * (u_light_attenuation.y + r * u_light_attenuation.z));

	out_color = vec4(color, 1.0);
}
)";


	static std::size_t bbox_to_screen_fan(geom::matrix<float, 4, 4> const & camera_transform, geom::box<float, 3> const & b, geom::point<float, 2> * result)
	{
		geom::point<float, 2> bbox_corners_screen[8];
		auto bbox_corners_screen_end = bbox_corners_screen;

		bool need_clipping = false;
		for (int z = 0; z < 2; ++z)
		{
			for (int y = 0; y < 2; ++y)
			{
				for (int x = 0; x < 2; ++x)
				{
					geom::point<float, 3> p;
					p[0] = geom::lerp<float>(b[0], x);
					p[1] = geom::lerp<float>(b[1], y);
					p[2] = geom::lerp<float>(b[2], z);

					auto q = camera_transform * geom::homogeneous(p);

					if (q[2] < -q[3] || q[2] > q[3])
					{
						need_clipping = true;
						break;
					}

					*bbox_corners_screen_end++ = {q[0] / q[3], q[1] / q[3]};
				}
			}
		}

		if (need_clipping)
		{
			// TODO: clip with near-Z manually instead of filling the entire screen
			*result++ = {-1.f, -1.f};
			*result++ = { 1.f, -1.f};
			*result++ = { 1.f,  1.f};
			*result++ = {-1.f,  1.f};
			return 4;
		}

		geom::point<float, 2> * bbox_hull_screen_it[8];
		auto bbox_hull_size = cg::graham_convex_hull(bbox_corners_screen, bbox_corners_screen_end, bbox_hull_screen_it) - bbox_hull_screen_it;

		for (std::size_t i = 0; i < bbox_hull_size; ++i)
			result[i] = *bbox_hull_screen_it[i];

		return bbox_hull_size;
	}

	struct deferred_renderer::impl
	{
		gfx::program g_buffer_pass_program{std::string(g_buffer_pass_common) + g_buffer_pass_vs, std::string(g_buffer_pass_common) + g_buffer_pass_fs};
		gfx::program ambient_pass_program{fullscreen_vs, ambient_pass_fs};
		gfx::program directional_light_pass_program{screen_vs, directional_light_pass_fs};
		gfx::program point_light_pass_program{screen_vs, point_light_pass_fs};

		// G-buffer attachments:
		// 0 - position (rbg)
		// 1 - albedo (rgb), lit (a)
		// 2 - normal (rgb)
		// 3 - material.diffuse (r), material.specular (g), material.shininess (b)

		gfx::framebuffer g_framebuffer;
		gfx::texture_2d g_buffer_texture[4];
		gfx::texture_2d g_buffer_depth;

		std::optional<geom::vector<std::size_t, 2>> g_buffer_size;

		gfx::mesh screen_mesh;
	};

	deferred_renderer::deferred_renderer()
		: pimpl_{std::make_unique<struct impl>()}
	{
		impl().g_buffer_pass_program.bind();
		impl().g_buffer_pass_program["u_texture"] = 0;

		for (std::size_t i = 0; i < 4; ++i)
		{
			impl().g_buffer_texture[i].nearest_filter();
		}

		impl().ambient_pass_program.bind();
		impl().ambient_pass_program["u_g0"] = 0;
		impl().ambient_pass_program["u_g1"] = 1;
		impl().ambient_pass_program["u_g2"] = 2;
		impl().ambient_pass_program["u_g3"] = 3;

		impl().directional_light_pass_program.bind();
		impl().directional_light_pass_program["u_g0"] = 0;
		impl().directional_light_pass_program["u_g1"] = 1;
		impl().directional_light_pass_program["u_g2"] = 2;
		impl().directional_light_pass_program["u_g3"] = 3;

		impl().point_light_pass_program.bind();
		impl().point_light_pass_program["u_g0"] = 0;
		impl().point_light_pass_program["u_g1"] = 1;
		impl().point_light_pass_program["u_g2"] = 2;
		impl().point_light_pass_program["u_g3"] = 3;

		impl().screen_mesh.setup<geom::point<float, 2>>();
	}

	deferred_renderer::~deferred_renderer() = default;

	static std::uint32_t const O_UNIFORM_COLOR  = 1 << 0;
	static std::uint32_t const O_TEXTURE_COLOR  = 1 << 1;
	static std::uint32_t const O_TRANSPARENT    = 1 << 2;
	static std::uint32_t const O_LIT            = 1 << 3;
	static std::uint32_t const O_CASTS_SHADOW   = 1 << 4;
	static std::uint32_t const O_PRE_TRANSFORM  = 1 << 5;
	static std::uint32_t const O_POST_TRANSFORM = 1 << 6;
	static std::uint32_t const O_INSTANCED      = 1 << 7;

	std::uint32_t mask(deferred_renderer::object const & o)
	{
		std::uint32_t m = 0;
		if (o.mat.color) m |= O_UNIFORM_COLOR;
		if (o.mat.texture) m |= O_TEXTURE_COLOR;
		if (o.mat.transparent) m |= O_TRANSPARENT;
		if (o.mat.lit) m |= O_LIT;
		if (o.casts_shadow) m |= O_CASTS_SHADOW;
		if (o.pre_transform) m |= O_PRE_TRANSFORM;
		if (o.post_transform) m |= O_POST_TRANSFORM;
		if (o.mesh->is_instanced()) m |= O_INSTANCED;
		return m;
	}

	void deferred_renderer::render(std::vector<object> const & objects, options const & opts)
	{
		// Get camera info

		assert(opts.camera);
		auto const camera_transform = opts.camera->transform();
		auto const camera_position = opts.camera->position();

		// Sort objects by mask & compute bbox

		std::unordered_map<std::uint32_t, std::vector<std::size_t>> objects_by_mask;

		geom::box<float, 3> lit_bbox;

		for (std::size_t i = 0; i < objects.size(); ++i)
		{
			auto const & o = objects[i];
			assert(o.mesh);

			if (o.mat.transparent) throw std::runtime_error("Transparency is not supported yet");

			objects_by_mask[mask(objects[i])].push_back(i);

			if (o.mat.lit) lit_bbox |= o.bbox;
		}

		// TODO: frustum culling

		// Resize g-buffer if needed

		auto buffer_size = geom::cast<std::size_t>(opts.viewport.dimensions());
		if (!impl().g_buffer_size || *impl().g_buffer_size != buffer_size)
		{
			// TODO: compact normals storage

			impl().g_buffer_texture[0].load<geom::vector<gfx::float16, 3>>(buffer_size);
			impl().g_buffer_texture[1].load<geom::vector<std::uint16_t, 4>>(buffer_size);
			impl().g_buffer_texture[2].load<geom::vector<gfx::float16, 3>>(buffer_size);
			impl().g_buffer_texture[3].load<geom::vector<gfx::float16, 3>>(buffer_size);
			impl().g_buffer_depth.load<gfx::depth24_pixel>(buffer_size);

			if (!impl().g_buffer_size)
			{
				for (std::size_t i = 0; i < 4; ++i)
				{
					impl().g_framebuffer.color(impl().g_buffer_texture[i], i);
				}
				impl().g_framebuffer.depth(impl().g_buffer_depth);
			}

			impl().g_framebuffer.assert_complete();

			impl().g_buffer_size = buffer_size;
		}

		// Setup g-buffer

		impl().g_framebuffer.bind();

		gl::Viewport(0, 0, opts.viewport[0].length(), opts.viewport[1].length());

		GLenum draw_buffers[4] { gl::COLOR_ATTACHMENT0, gl::COLOR_ATTACHMENT1, gl::COLOR_ATTACHMENT2, gl::COLOR_ATTACHMENT3 };
		gl::DrawBuffers(4, draw_buffers);
		check_error();
		float buffer_1_clear[4] { 0.f, 0.f, 0.f, 0.f };
		if (opts.clear_color)
		{
			buffer_1_clear[0] = (*opts.clear_color)[0] / opts.max_intensity;
			buffer_1_clear[1] = (*opts.clear_color)[1] / opts.max_intensity;
			buffer_1_clear[2] = (*opts.clear_color)[2] / opts.max_intensity;
		}

		gl::ClearBufferfv(gl::COLOR, 1, buffer_1_clear);

		gl::ClearDepth(1.f);
		gl::Clear(gl::DEPTH_BUFFER_BIT);

		gl::Enable(gl::DEPTH_TEST);
		gl::DepthFunc(gl::LEQUAL);

		gl::Disable(gl::BLEND);

		gl::Enable(gl::CULL_FACE);
		gl::CullFace(gl::BACK);

		// Render to g-buffer

		impl().g_buffer_pass_program.bind();
		impl().g_buffer_pass_program["u_camera_transform"] = camera_transform;
		impl().g_buffer_pass_program["u_max_intensity"] = opts.max_intensity;

		for (auto const & p : objects_by_mask)
		{
			if (p.second.empty()) continue;

			std::uint32_t mask = p.first;

			impl().g_buffer_pass_program["u_flag_mask"] = mask;

			for (std::size_t i : p.second)
			{
				auto const & o = objects[i];

				if (mask & O_UNIFORM_COLOR)
					impl().g_buffer_pass_program["u_color"] = *o.mat.color;

				if (mask & O_TEXTURE_COLOR)
				{
					gl::ActiveTexture(gl::TEXTURE0);
					o.mat.texture->bind();
				}

				if (mask & O_PRE_TRANSFORM)
					impl().g_buffer_pass_program["u_pre_transform"] = *o.pre_transform;

				if (mask & O_POST_TRANSFORM)
					impl().g_buffer_pass_program["u_post_transform"] = *o.post_transform;

				impl().g_buffer_pass_program["u_material"] = geom::vector<float, 3>{o.mat.diffuse, o.mat.specular.intensity, o.mat.specular.shininess};

				o.mesh->draw();
			}
		}

		// Setup destination framebuffer

		assert(opts.framebuffer);
		opts.framebuffer->bind();

		gl::DrawBuffer(opts.draw_buffer);

		gl::Viewport(opts.viewport[0].min, opts.viewport[1].min, opts.viewport[0].length(), opts.viewport[1].length());

		gl::Disable(gl::DEPTH_TEST);
		gl::Disable(gl::BLEND);

		gl::ActiveTexture(gl::TEXTURE0);
		impl().g_buffer_texture[0].bind();
		gl::ActiveTexture(gl::TEXTURE1);
		impl().g_buffer_texture[1].bind();
		gl::ActiveTexture(gl::TEXTURE2);
		impl().g_buffer_texture[2].bind();
		gl::ActiveTexture(gl::TEXTURE3);
		impl().g_buffer_texture[3].bind();

		impl().screen_mesh.bind();

		// TODO: directional light shadows

		// TODO: point light shadows

		// Draw unlit & ambient layers

		impl().ambient_pass_program.bind();
		impl().ambient_pass_program["u_ambient"] = opts.ambient;
		impl().ambient_pass_program["u_max_intensity"] = opts.max_intensity;
		gl::DrawArrays(gl::TRIANGLES, 0, 6);

		gl::Enable(gl::BLEND);
		gl::BlendFunc(gl::ONE, gl::ONE);

		// Directional lights

		geom::point<float, 2> bbox_hull_screen[8];
		auto lit_bbox_hull_size = bbox_to_screen_fan(camera_transform, lit_bbox, bbox_hull_screen);
		impl().screen_mesh.load(bbox_hull_screen, lit_bbox_hull_size, gl::TRIANGLE_FAN);

		impl().directional_light_pass_program.bind();
		impl().directional_light_pass_program["u_camera_position"] = camera_position;
		impl().directional_light_pass_program["u_max_intensity"] = opts.max_intensity;

		for (auto const & l : opts.directional_lights)
		{
			impl().directional_light_pass_program["u_light_direction"] = geom::normalized(l.direction);
			impl().directional_light_pass_program["u_light_color"] = l.color;
			impl().screen_mesh.draw();
		}

		// Point lights

		float min_intensity = opts.min_intensity.value_or(opts.max_intensity / 256.f);

		impl().point_light_pass_program.bind();
		impl().point_light_pass_program["u_camera_position"] = camera_position;
		impl().point_light_pass_program["u_max_intensity"] = opts.max_intensity;

		for (auto const & l : opts.point_lights)
		{
			float I = std::max({l.color[0], l.color[1], l.color[2]});
			auto r = geom::solve_quadratic(l.attenuation.c2, l.attenuation.c1, l.attenuation.c0 - I / min_intensity);
			geom::box<float, 3> light_bbox;
			if (r)
			{
				float light_influence_radius = r->second;

				for (std::size_t i = 0; i < 3; ++i)
					light_bbox[i] = {l.position[i] - light_influence_radius, l.position[i] + light_influence_radius};

				light_bbox &= lit_bbox;
			}
			else
			{
				light_bbox = lit_bbox;
			}

			auto light_bbox_hull_size = bbox_to_screen_fan(camera_transform, light_bbox, bbox_hull_screen);
			impl().screen_mesh.load(bbox_hull_screen, light_bbox_hull_size, gl::TRIANGLE_FAN);

			impl().point_light_pass_program["u_light_position"] = l.position;
			impl().point_light_pass_program["u_light_color"] = l.color;
			impl().point_light_pass_program["u_light_attenuation"] = geom::vector{l.attenuation.c0, l.attenuation.c1, l.attenuation.c2};
			impl().screen_mesh.draw();
		}

		check_error();
	}

}