#include "rendering/batch_renderers/point_cloud_batch_renderer.hpp"

#include <algorithm>

#include "rendering/requirements/point_cloud_requirements.hpp"
#include "shading/uniforms/point_cloud_uniforms.hpp"
#include "util/unroll_bool_template.hpp"

namespace rendering
{

point_cloud_batch_renderer::point_cloud_batch_renderer(int render_mode_count)
	: m_render_mode_count{ render_mode_count } {};

std::pair<std::size_t, bool> point_cloud_batch_renderer::lookup_batch(
	const batch_components_type& batch_component
) const {
	const auto component_it = std::upper_bound(
		m_component_lookup.begin(), m_component_lookup.end(),
		batch_component,
		[](const auto& batch_component, const auto& entry)
		{
			return batch_component < entry.first;
		}
	);

	const auto index = component_it - m_component_lookup.begin();

	const auto match = (
		index == 0 or m_component_lookup[index - 1].first == batch_component
	);

	return { index - static_cast<std::size_t>(match), match };
}

std::optional<point_cloud_batch_renderer::id_type> point_cloud_batch_renderer::add(
	batch_components_type batch_components,
	const zgl::point_cloud_handle& point_cloud,
	const aabb& bounding_box,
	const zgl::model_matrix_handle& transform,
	const shader_program_lookups::point_cloud_lookup& shader_program_lookup
) {
	const auto [ lookup_index, lookup_match ] = lookup_batch(batch_components);

	std::size_t batch_index;
	batch_id_type batch_id;

	if (lookup_match) {
		batch_index = m_component_lookup[lookup_index].second;
		batch_id = m_id_lookup[batch_index];
	}
	else
	{
		auto base_requirements = requirements::point_cloud::flags{};

		const auto vertex_comps = batch_components;

		for (std::size_t i{}; i != requirements::point_cloud::all.size(); ++i)
		{
			const auto& requirement = requirements::point_cloud::all[i];

			if (
				(vertex_comps & requirement.vertex_requirements) != components::point_cloud_vertex::flags::none
			) {
				base_requirements |= requirements::point_cloud::flags{ 1 << i };
			}
		}

		const auto uniform_color_shader = shader_program_lookup.find(
			base_requirements | requirements::point_cloud::flags::uniform_color
		);

		if (not uniform_color_shader)
		{
			return std::nullopt;
		}

		const auto rainbow_shader = shader_program_lookup.find(
			base_requirements | requirements::point_cloud::flags::rainbow
		);

		if (not rainbow_shader)
		{
			return std::nullopt;
		}

		auto shader_programs = std::array{
			*uniform_color_shader,
			*rainbow_shader
		};

		batch_index = m_batches.size();
		batch_id = m_next_batch_id++;
		m_batches.emplace_back(batch_type{}, batch_components);
		m_id_lookup.push_back(batch_id);
		m_component_lookup.emplace(m_component_lookup.begin() + lookup_index, batch_components, batch_index);
		m_shader_programs.insert(
			m_shader_programs.begin() + lookup_index,
			shader_programs.begin(), shader_programs.end()
		);
	}

	auto& batch = m_batches[batch_index].first;

	const auto mesh_id = batch.add(point_cloud, bounding_box, transform);

	return id_type{ batch_id, mesh_id };
}

std::optional<aabb> point_cloud_batch_renderer::bounding_box(id_type id)
{
	const auto lookup_it = std::ranges::find(m_id_lookup, id.first);

	if (lookup_it == m_id_lookup.end())
	{
		return std::nullopt;
	}

	const auto batch_index = lookup_it - m_id_lookup.begin();
	auto& batch = m_batches[batch_index].first;

	return batch.bounding_box(id.second);
}

bool point_cloud_batch_renderer::remove(const id_type id)
{
	const auto lookup_it = std::ranges::find(m_id_lookup, id.first);

	if (lookup_it == m_id_lookup.end())
	{
		return false;
	}

	const auto batch_index = lookup_it - m_id_lookup.begin();
	auto& batch = m_batches[batch_index].first;

	// If batches can be removed the indices in m_component_lookup need to be changed.
	return batch.remove(id.second);
}


template<bool Normals>
void render_point_cloud_batch(
	const zgl::shader_program_handle& shader_program,
	const point_cloud_batch& batch,
	const glm::mat4& vp_matrix,
	const glm::vec3& camera_position
) {
	const auto point_clouds = batch.point_clouds();
	const auto transforms = batch.transforms();

	namespace uniforms = shading::uniforms::point_cloud;

	for (std::size_t i{}; i != point_clouds.size(); ++i)
	{
		const auto& point_cloud = point_clouds[i];
		const auto& model_matrix = transforms[i];

		// TODO check order
		const auto mvp_matrix = vp_matrix * model_matrix;

		shader_program.set_uniform<uniforms::mvp.info>(mvp_matrix);

		if constexpr (Normals)
		{
			shader_program.set_uniform<uniforms::model.info>(model_matrix);
			shader_program.set_uniform<uniforms::camera_position.info>(camera_position);
		}

		point_cloud.bind();

		using block_index_type = ztu::u16;
		static constexpr auto block_size = static_cast<std::size_t>(
			std::numeric_limits<block_index_type>::max()
		);

		for (GLsizei j{}; j < point_cloud.point_count; j += block_size)
		{
			const auto points_left = static_cast<std::size_t>(point_cloud.point_count) - j;
			const auto points_in_block = std::min(points_left, block_size);
			glDrawArrays(
				GL_POINTS,
				j,
				static_cast<block_index_type>(points_in_block)
			);
		}
	}
}

void point_cloud_batch_renderer::render(
	const modes::point_cloud render_mode,
	const glm::mat4& vp_matrix,
	const glm::vec3& camera_position,
	const lighting_setup&
) {

	namespace uniforms = shading::uniforms::point_cloud;

	const auto render_mode_index = static_cast<std::size_t>(render_mode);

	glEnable(GL_PROGRAM_POINT_SIZE);
	glEnable(GL_POINT_SMOOTH);

	const auto rainbow = render_mode == modes::point_cloud::rainbow;


	for (std::size_t i{}; i != m_batches.size(); ++i) {

		const auto& [ batch, vertex_components ] = m_batches[i];

		const auto normals = static_cast<bool>(vertex_components & components::point_cloud_vertex::flags::normal);

		const auto& shader_program = m_shader_programs[i * m_render_mode_count + render_mode_index];

		shader_program.bind();

		if (rainbow)
		{
			shader_program.set_uniform<uniforms::rainbow_offset_y.info>(0.0f); // TODO fix
			shader_program.set_uniform<uniforms::rainbow_scale_y.info>(0.0f); // TODO fix
		}

		unroll_bool_function_template(
			[&]<bool Normals>() {
				render_point_cloud_batch<Normals>(
					shader_program,
					batch,
					vp_matrix,
					camera_position
				);
			},
			normals
		);
	}

	glDisable(GL_PROGRAM_POINT_SIZE);
	glDisable(GL_POINT_SMOOTH);

	zgl::point_cloud_handle::unbind();
	zgl::shader_program_handle::unbind();
}

}