#include "rendering/batch_renderers/mesh_batch_renderer.hpp"

#include "shader_program/uniforms/mesh_uniforms.hpp"
#include "util/unroll_bool_template.hpp"
#include "util/logger.hpp" // TODO remove
#include <bitset> // TODOE remove

namespace rendering
{

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

std::pair<std::size_t, bool> mesh_batch_renderer::lookup_batch(
	const batch_components_type& batch_components
) const {

	const auto component_it = std::upper_bound(
		m_component_lookup.begin(), m_component_lookup.end(),
		batch_components,
		[](const auto& components, const auto& entry)
		{
			return components < entry.first;
		}
	);

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

	const auto match = (
		index != 0 and m_component_lookup[index - 1].first == batch_components
	);

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

std::optional<mesh_batch_renderer::id_type> mesh_batch_renderer::add(
	const batch_components_type& batch_component,
	const zgl::mesh_handle& mesh,
	const aabb& bounding_box,
	const zgl::model_matrix_handle& transform,
	const zgl::material_handle& material,
	const shader_program_lookups::mesh_lookup& shader_program_lookup
) {
	const auto [ lookup_index, lookup_match ] = lookup_batch(batch_component);

	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::mesh::flags::position;

		const auto [ vertex_comps, material_comps ] = batch_component;

		// If no texture is provided, the uniform color is provided by ambient light.
		if ((material_comps & material_component::flags::texture) == material_component::flags::none)
		{
			base_requirements |= requirements::mesh::flags::uniform_color;
		}

		ztu::logger::debug("vertex_comps: %", std::bitset<32>{ static_cast<unsigned long long>(static_cast<int>(vertex_comps)) });
		ztu::logger::debug("material_comps: %", std::bitset<32>{ static_cast<unsigned long long>(static_cast<int>(material_comps)) });
		ztu::logger::debug("lit reqs: %", std::bitset<32>{ static_cast<unsigned long long>(static_cast<int>(shader_program::features::mesh::lit.uniforms)) });


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

			if (
				(
					requirement.vertex_requirements != components::mesh_vertex::flags::none and
					(vertex_comps & requirement.vertex_requirements) == requirement.vertex_requirements
				)
				and
				(
					requirement.material_requirements != material_component::flags::none and
					(material_comps & requirement.material_requirements) == requirement.material_requirements
				)
			) {
				base_requirements |= requirements::mesh::flags{ 1 << i };
			}
		}

		ztu::logger::debug("base reqs: %", std::bitset<32>{ static_cast<unsigned long long>(static_cast<int>(base_requirements)) });


		const auto base_shader = shader_program_lookup.find(base_requirements);
		if (not base_shader)
		{
			ztu::logger::warn("Could not find base shader!");
			return std::nullopt;
		}

		const auto point_shader = shader_program_lookup.find(base_requirements | requirements::mesh::flags::point);
		if (not point_shader)
		{
			ztu::logger::warn("Could not find point shader!");
			return std::nullopt;
		}

		const auto lit_shader = shader_program_lookup.find(base_requirements | requirements::mesh::flags::lit);
		if (not lit_shader)
		{
			ztu::logger::warn("Could not find lit shader!");
			return std::nullopt;
		}

		auto shader_programs = std::array{
			*base_shader,
			*point_shader,
			*base_shader,
			*lit_shader
		};

		ztu::logger::debug(
			"shaders: % % %",
			base_shader->program_id,
			point_shader->program_id,
			lit_shader->program_id
		);

		batch_index = m_batches.size();
		batch_id = m_next_batch_id++;
		m_batches.emplace_back(batch_type{}, batch_component);
		m_id_lookup.push_back(batch_id);
		m_component_lookup.emplace(m_component_lookup.begin() + lookup_index, batch_component, 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(mesh, bounding_box, transform, material);

	return id_type{ batch_id, mesh_id };
}

std::optional<aabb> mesh_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 mesh_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 Textured, bool Lit, bool Alpha>
void render_mesh_batch(
	const zgl::shader_program_handle& shader_program,
	const mesh_batch& batch,
	const glm::mat4& vp_matrix,
	const glm::mat4& view_matrix,
	const GLenum draw_mode
) {
	const auto meshes = batch.meshes();
	const auto transforms = batch.transforms();
	const auto textures = batch.textures();
	const auto surface_properties = batch.surface_properties();
	const auto alphas = batch.alphas();

	/*ztu::logger::debug("meshes: %", meshes.size());
	ztu::logger::debug("transforms: %", transforms.size());
	ztu::logger::debug("textures: %", textures.size());
	ztu::logger::debug("surface_properties: %", surface_properties.size());
	ztu::logger::debug("alphas: %", alphas.size());

	ztu::logger::debug("textured: % alpha: % lit: %", Textured, Alpha, Lit);*/

	namespace uniforms = shader_program::uniforms::mesh;

	for (std::size_t i{}; i != meshes.size(); ++i)
	{
		//ztu::logger::debug("Mesh: %", i);

		const auto& mesh = meshes[i];
		const auto& model_matrix = transforms[i];

		const auto mvp_matrix = vp_matrix * model_matrix;

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

		if constexpr (Textured)
		{
			textures[i].bind();
		}

		if constexpr (Lit)
		{
			shader_program.set_uniform<uniforms::model_matrix.info>(model_matrix);

			// TODO more efficient set
			const auto& properties = surface_properties[i];
			shader_program.set_uniform<uniforms::ambient_filter.info>(properties.ambient_filter);
			shader_program.set_uniform<uniforms::diffuse_filter.info>(properties.diffuse_filter);
			shader_program.set_uniform<uniforms::specular_filter.info>(properties.specular_filter);
			shader_program.set_uniform<uniforms::shininess.info>(properties.shininess);
		}

		if constexpr (Alpha)
		{

			shader_program.set_uniform<uniforms::alpha.info>(alphas[i]);
		}

		//ztu::logger::debug("vao: % valid: %%", mesh.vao_id, std::boolalpha, (bool)glIsVertexArray(mesh.vao_id));

		mesh.bind();

		//ztu::logger::debug("glDrawElements(%, %)", draw_mode, mesh.index_count);

		glDrawElements(draw_mode, mesh.index_count, GL_UNSIGNED_INT, nullptr);

		//ztu::logger::debug("done");
	}
}

void mesh_batch_renderer::render(
	const modes::mesh render_mode,
	const glm::mat4& vp_matrix,
	const glm::mat4& view_matrix,
	const glm::vec3& view_pos,
	const lighting_setup& lights
) {

	namespace uniforms = shader_program::uniforms::mesh;

	const auto render_mode_index = static_cast<int>(render_mode);

	const auto lines = render_mode == modes::mesh::wire_frame;
	const auto points = render_mode == modes::mesh::points;
	const auto lit = render_mode == modes::mesh::lit_faces;

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

		//ztu::logger::debug("batch: %", i);

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

		const auto [ vertex_components, material_components ] = batch_components;

		const auto textured = (
			(vertex_components & components::mesh_vertex::flags::tex_coord) != components::mesh_vertex::flags::none
			and (material_components & material_component::flags::texture) !=  material_component::flags::none
		);

		const auto alpha = (
			(material_components & material_component::flags::transparency) != material_component::flags::none
		);

		const auto draw_mode = points ? GLenum{ GL_POINTS } : GLenum{ GL_TRIANGLES };

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

		//ztu::logger::debug("shader_program: % valid: %%", shader_program.program_id, std::boolalpha, (bool)glIsProgram(shader_program.program_id));

		shader_program.bind();

		if (lit)
		{
			// TODO set more efficiently
			shader_program.set_uniform<uniforms::view_pos.info>(view_pos);
			shader_program.set_uniform<uniforms::point_light_direction.info>(lights.point_light_direction);
			shader_program.set_uniform<uniforms::point_light_color.info>(lights.point_light_color);
			shader_program.set_uniform<uniforms::ambient_light_color.info>(lights.ambient_light_color);
		}

		if (textured)
		{
			constexpr auto texture_unit = 0;
			glActiveTexture(GL_TEXTURE0 + texture_unit);
			shader_program.set_uniform<uniforms::tex.info>(texture_unit);
		}
		else
		{
			shader_program.set_uniform<uniforms::color.info>(glm::vec4(lights.ambient_light_color, 1.0f));
		}

		if (lines)
		{
			glPolygonMode(GL_FRONT, GL_LINE);
			glPolygonMode(GL_BACK, GL_LINE);
		}

		unroll_bool_function_template(
			[&]<bool Textured, bool Lit, bool Alpha>() {
				render_mesh_batch<Textured, Lit, Alpha>(
					shader_program,
					batch,
					vp_matrix,
					view_matrix,
					draw_mode
				);
			},
			textured, lit, alpha
		);

		if (lines)
		{
			glPolygonMode(GL_FRONT, GL_FILL);
			glPolygonMode(GL_BACK, GL_FILL);
		}
	}

	zgl::texture_handle::unbind();
	zgl::mesh_handle::unbind();
	zgl::shader_program_handle::unbind();
}

}