#include "opengl/data_managers/shader_program_manager.hpp"


struct prioritized_metadata_comparator
{
	using type = zgl::shader_program_metadata;

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

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

void zgl::shader_program_manager::process(
	const assets::shader_source_store& shader_sources
) {
	m_shader_manager.preprocess(shader_sources);
}

void zgl::shader_program_manager::get_handles(
	const assets::shader_source_store& shader_sources,
	std::span<const shading::shader_program_requirements> requirements,
	std::span<shader_program_metadata> metadata,
	std::span<shader_program_handle> shader_programs
) {
	m_shader_requirements_buffer.clear();

	for (auto [ req, program_meta, program_handle ] : std::ranges::views::zip(
		requirements,
		metadata,
		shader_programs
	)) {

		if (auto shader_match = find_shader_program(req))
		{
			const auto& [ meta, handle ] = *shader_match;
			program_meta.static_enabled = meta.static_enabled;
			program_meta.dynamic_enable = meta.dynamic_enable;
			program_handle = handle;
		}
		else
		{
			program_meta = {};
			program_handle = {};
			m_shader_requirements_buffer.emplace_back(
				req.geometry,
				req.features
			);
		}
	}


	m_shader_metadata_buffer.clear();
	m_shader_metadata_buffer.resize(m_shader_requirements_buffer.size());
	shader_set_buffer.clear();

	m_shader_manager.get_handles(
		shader_sources,
		m_shader_requirements_buffer,
		m_shader_metadata_buffer,
		shader_set_buffer
	);

	auto shader_set_req_it = m_shader_requirements_buffer.begin();
	auto shader_set_meta_it = m_shader_metadata_buffer.begin();
	auto shader_set_it = shader_set_buffer.begin();

	const auto prev_shader_program_count = m_shader_program_lookup.size();

	for (auto [ program_meta, program_handle ] : std::ranges::views::zip( metadata, shader_programs))
	{
		if (not program_handle.valid())
		{
			if (std::ranges::any_of(shader_set_it->stages, &shader_handle::valid))
			{
				shader_program_data program{};

				if (link_shader_program(*shader_set_it))
				{
					program_handle = program.handle;

					program_meta = shader_program_metadata{
						.geometry = shader_set_req_it->geometry,
						.static_enabled = shader_set_meta_it->static_enabled,
						.dynamic_enable = shader_set_meta_it->dynamic_enable
					};

					m_shader_program_lookup.emplace_back(program_meta, std::move(program));
			   }
			}

			++shader_set_req_it;
			++shader_set_meta_it;
			++shader_set_it;
		}

	}

	const auto new_shader_programs = std::span(m_shader_program_lookup).subspan(prev_shader_program_count);

	std::ranges::sort(
		new_shader_programs,
		prioritized_metadata_comparator{},
		&shader_program_lookup_entry_type::first
	);

	std::ranges::inplace_merge(
		m_shader_program_lookup,
		m_shader_program_lookup.begin() + prev_shader_program_count,
		prioritized_metadata_comparator{},
		&shader_program_lookup_entry_type::first
	);
}

std::optional<std::pair<zgl::shader_program_metadata, zgl::shader_program_handle>> zgl::shader_program_manager::find_shader_program(
	const shading::shader_program_requirements& requirements
) {

	auto shader_program_it = std::ranges::lower_bound(
		m_shader_program_lookup,
		requirements.geometry,
		std::greater{},
		[](const shader_program_lookup_entry_type& entry)
		{
			const auto& meta = entry.first;
			return meta.geometry;
		}
	);

	while (
		shader_program_it != m_shader_program_lookup.end() and
		shader_program_it->first.geometry == requirements.geometry
	) {
		const auto& [ meta, data ] = *shader_program_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_program_it;
	}

	return std::nullopt;
}

bool zgl::shader_program_manager::link_shader_program(
	const shader_handle_set& shaders
) {
	const auto program = shader_program_data{ glCreateProgram() };
	if (const auto e = get_error())
	{
		ztu::logger::error("Error while creating shader program: %.", e.message());
		return false;
	}

	for (const auto [ index, entry ] : std::ranges::views::enumerate(std::ranges::views::zip(
		shaders.stages,
		shading::stage::names
	))) {
		const auto& [ shader, name ] = entry;

		auto attached = false;

		if (shader.id)
		{
			glAttachShader(program.handle.id, shader.id);
			if (const auto e = get_error())
			{
				if (shading::stage::mandatory[index])
				{
					ztu::logger::error("Error while attaching the mandatory % shader: %.", name, e.message());
					return false;
				}
				ztu::logger::warn("Error while attaching the optional % shader: %.", name, e.message());
			}
			else
			{
				attached = true;
			}
		}

		if (not attached)
		{
			ztu::logger::warn("Using default % shader.", name);
		}
	}

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

	auto status = GLint{ GL_FALSE };
	glGetProgramiv(program.handle.id, GL_LINK_STATUS, &status);
	if (const auto e = get_error())
		{
		ztu::logger::error("Error while retrieving shader program link status: %.", e.message());
		return false;
	}

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

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

		ztu::logger::error("Error while linking program:\n%", log);

		return false;
	}

	glUseProgram(0);
	if (const auto e = get_error())
	{
		ztu::logger::warn("Error while resetting active shader program: %.", e.message());
	}

	for (const auto& shader : shaders.stages)
	{
		if (shader.id)
		{
			glDetachShader(program.handle.id, shader.id);
			if (const auto e = get_error())
			{
				ztu::logger::warn("Error while detaching shader: %.", e.message());
			}
		}
	}

	return true;
}