#include "opengl/data_uploaders/shader_compiler.hpp"

#include <numeric>
#include <bits/ranges_algobase.h>
#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
		);
	}
};


zgl::shader_handle zgl::shader_program_compiler::find_shader(
	const shading::shader_requirements& requirements
) {

	auto handle = shader_handle{};

	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 };
		}
	);

	dynamic_shader_source_store::id_type source_id{};
	shading::features::generic::type to_be_enabled{};

	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)
		{
			to_be_enabled = req.features & static_enable;
			source_id = id;

			res.static_enabled = features & ~dynamic_enable & ~unwanted_features;
			res.dynamic_enable = dynamic_enable;

			break;
		}

		++shader_it;
	}

	if (source_id)
	{
		const auto [ shader_source_it, source_found ] = shader_sources.find(source_id);
		if (source_found)
		{
			get_define_strings(
				req.geometry,
				to_be_enabled,
				res.string_count,
				shader_strings
			);
		}
	}

	return handle;
}

void zgl::shader_program_compiler::compile_shaders(
	const dynamic_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);
			handle = find_shader(shader_req);
			if (not handle.valid())
			{
				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.fetch_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))
	{
		for (auto [ stage_index, handle ] : std::ranges::views::enumerate(shader_set.stages))
		{
			if (not handle.valid())
			{
				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));
				   }
				}

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

	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_program_compiler::compile_shader(
	GLenum shader_type,
	std::span<const char*> source_strings,
	shader_data& shader
) {

	shader = shader_data(glCreateShader(shader_type));

	glShaderSource(
		shader.handle.id,
		static_cast<GLsizei>(source_strings.size()),
		source_strings.data(),
		nullptr
	);

	glCompileShader(shader.handle.id);

	GLint success;
	glGetShaderiv(shader.handle.id, GL_COMPILE_STATUS, &success);

	if (not success)
	{
		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;
}

void zgl::shader_program_compiler::preprocess(
	const dynamic_shader_source_store& shader_sources
) {
	m_preprocessor.preprocess(shader_sources);
}