#include "opengl/data_managers/shader_manager.hpp"

#include <numeric>
#include <bits/ranges_algobase.h>

#include "opengl/error.hpp"
#include "util/logger.hpp"


struct prioritized_metadata_comparator
{
	using type = zgl::shader_metadata;

	bool operator()(const type& a, const type& b) const noexcept
	{
		if (a.geometry != b.geometry)
		{
			return a.geometry > b.geometry;
		}

		if (a.stage != b.stage)
		{
			return a.stage > b.stage;
		}

		static constexpr auto more_features = std::popcount<zgl::shading::features::generic::type>;

		return std::ranges::lexicographical_compare(
			std::array{ a.dynamic_enable, a.static_enabled },
			std::array{ b.dynamic_enable, b.static_enabled },
			std::greater{},
			more_features,
			more_features
		);
	}
};

static constexpr auto gl_shader_types = std::array<GLenum, zgl::shading::stage::count>{
	GL_VERTEX_SHADER,
	GL_TESS_CONTROL_SHADER,
	GL_TESS_EVALUATION_SHADER,
	GL_GEOMETRY_SHADER,
	GL_FRAGMENT_SHADER
};

std::optional<std::pair<zgl::shader_metadata, zgl::shader_handle>> zgl::shader_manager::find_shader(
	const shading::shader_requirements& requirements
) {
	auto shader_it = std::ranges::lower_bound(
		m_shader_lookup,
		std::pair{ requirements.geometry, requirements.stage },
		std::greater{},
		[](const shader_lookup_entry_type& entry)
		{
			const auto& meta = entry.first;
			return std::pair{ meta.geometry, meta.stage };
		}
	);

	while (
		shader_it != m_shader_lookup.end() and
		shader_it->first.geometry == requirements.geometry and
		shader_it->first.stage == requirements.stage
	) {
		const auto& [ meta, data ] = *shader_it;

		const auto unwanted_static_features = meta.static_enabled & ~requirements.features;
		const auto required_dynamic_features = requirements.features & ~meta.static_enabled;
		const auto missing_dynamic_features = required_dynamic_features & ~meta.dynamic_enable;

		if (unwanted_static_features == 0 and missing_dynamic_features == 0)
		{
			return std::pair{ meta, data.handle };
		}

		++shader_it;
	}

	return std::nullopt;
}

void zgl::shader_manager::process(const store_type& shader_sources)
{
	m_preprocessor.process(shader_sources);
}

void zgl::shader_manager::get_handles(
	const assets::shader_source_store& shader_sources,
	std::span<const shading::shader_set_requirements> requirements,
	std::span<shader_set_metadata> metadata,
	std::span<shader_handle_set> shader_sets
) {
	for (auto [ req, shader_set_meta, shader_set ] : std::ranges::views::zip(
		requirements,
		metadata,
		shader_sets
	)) {
		shader_set = {};
		shader_set_meta = {};

		auto shader_req = shading::shader_requirements{
			.geometry = req.geometry,
			.stage = {},
			.features = req.features
		};

		for (auto [ stage_index, handle ] : std::ranges::views::enumerate(shader_set.stages))
		{
			shader_req.stage = static_cast<shading::stage::types>(stage_index);

			if (auto shader_match = find_shader(shader_req))
			{
				const auto& [ shader_meta, shader_handle ] = *shader_match;
				shader_set_meta.static_enabled |= shader_meta.static_enabled;
				shader_set_meta.dynamic_enable |= shader_meta.dynamic_enable;
				handle = shader_handle;
			}
			else
			{
				m_source_requirement_buffer.push_back(shader_req);
			}
		}
	}


	m_preprocessed_shader_source_metadata_buffer.clear();
	m_preprocessed_shader_source_metadata_buffer.resize(m_source_requirement_buffer.size());
	m_source_strings_buffer.clear();

	m_preprocessor.get_shader_sources(
		shader_sources,
		m_source_requirement_buffer,
		m_preprocessed_shader_source_metadata_buffer,
		m_source_strings_buffer
	);

	auto source_strings_it = m_source_strings_buffer.begin();
	auto source_req_it = m_source_requirement_buffer.begin();
	auto source_meta_it = m_preprocessed_shader_source_metadata_buffer.begin();

	const auto prev_shader_count = m_shader_lookup.size();

	for (auto [ shader_set_meta, shader_set ] : std::ranges::views::zip( metadata, shader_sets))
	{
		auto shader_missing = false;

		for (auto [ stage_index, handle ] : std::ranges::views::enumerate(shader_set.stages))
		{
			if (not handle.valid())
			{
				if (not shader_missing)
				{
					if (source_meta_it->string_count > 0)
					{
						shader_data shader{};

						if (compile_shader(
						   gl_shader_types[stage_index],
						   std::span(source_strings_it.base(), source_meta_it->string_count),
						   shader
					   )) {
							handle = shader.handle;
							auto shader_meta = shader_metadata{
								.geometry = source_req_it->geometry,
								.stage = source_req_it->stage,
								.static_enabled = source_meta_it->static_enabled,
								.dynamic_enable = source_meta_it->dynamic_enable
							};

							shader_set_meta.static_enabled |= shader_meta.static_enabled;
							shader_set_meta.dynamic_enable |= shader_meta.dynamic_enable;

							m_shader_lookup.emplace_back(shader_meta, std::move(shader));
					   }
					}

					shader_missing = not handle.valid() and shading::stage::mandatory[stage_index];
				}

				source_strings_it += source_meta_it->string_count;
				++source_meta_it;
				++source_req_it;
			}
		}

		if (shader_missing)
		{
			shader_set.stages = {};
		}
	}

	const auto new_shaders = std::span(m_shader_lookup).subspan(prev_shader_count);

	std::ranges::sort(
		new_shaders,
		prioritized_metadata_comparator{},
		&shader_lookup_entry_type::first
	);

	std::ranges::inplace_merge(
		m_shader_lookup,
		m_shader_lookup.begin() + prev_shader_count,
		prioritized_metadata_comparator{},
		&shader_lookup_entry_type::first
	);
}

bool zgl::shader_manager::compile_shader(
	GLenum shader_type,
	std::span<const char*> source_strings,
	shader_data& shader
) {

	shader = shader_data(glCreateShader(shader_type));
	if (const auto e = get_error())
	{
		ztu::logger::error("Error while creating shader: %", e.message());
		return false;
	}

	glShaderSource(
		shader.handle.id,
		static_cast<GLsizei>(source_strings.size()),
		source_strings.data(),
		nullptr
	);
	if (const auto e = get_error())
	{
		ztu::logger::error("Error while setting shader source: %", e.message());
		return false;
	}

	glCompileShader(shader.handle.id);
	if (const auto e = get_error())
	{
		ztu::logger::error("Error while compiling shader: %", e.message());
		return false;
	}

	auto status = GLint{ GL_FALSE };
	glGetShaderiv(shader.handle.id, GL_COMPILE_STATUS, &status);

	if (const auto e = get_error())
	{
		ztu::logger::error("Error while retrieving shader compilation status: %", e.message());
		return false;
	}

	if (status == GL_FALSE)
	{
		GLint log_length{};
		glGetShaderiv(shader.handle.id, GL_INFO_LOG_LENGTH, &log_length);

		auto log = std::string(log_length, ' ');
		glGetShaderInfoLog(shader.handle.id, log_length, nullptr, log.data());

		ztu::logger::error("Error while compiling shader:\n%", log);

		return false;
	}

	return true;
}