#include "opengl/shader_program_lookup.hpp"

#include <algorithm>
#include <bitset>

#include "util/logger.hpp" // TODO remove

namespace zgl
{


void shader_program_lookup::add(
	const shader_program_handle& shader_program_handle,
	const std::span<const shader_program_variable> all_attributes,
	const std::span<const shader_program_variable> all_uniforms
) {
	const auto attributes = shader_program_handle.check_attribute_support(all_attributes);
	const auto uniforms = shader_program_handle.check_uniform_support(all_uniforms);

	ztu::logger::debug("add [%] uniforms: % attributes: %",
			shader_program_handle.program_id,
			std::bitset<32>(uniforms),
			std::bitset<32>(attributes)
		);

	const auto lower_uniform = std::ranges::lower_bound(
		m_mesh_shader_program_uniforms,
		uniforms
	);

	const auto upper_uniform = std::find_if(
		lower_uniform, m_mesh_shader_program_uniforms.end(),
		[&](const auto& curr_uniforms) {
			return curr_uniforms > uniforms;
		}
	);

	const auto lower_index = lower_uniform - m_mesh_shader_program_uniforms.begin();
	const auto upper_index = upper_uniform - m_mesh_shader_program_uniforms.begin();

	const auto lower_attribute = m_mesh_shader_program_attributes.begin() + lower_index;
	const auto upper_attribute = m_mesh_shader_program_attributes.begin() + upper_index;

	const auto attribute_it = std::upper_bound(
		lower_attribute, upper_attribute,
		attributes,
		[](const auto& attributes, const auto& entry) {
			return attributes < entry.attributes;
		}
	);

	const auto index = attribute_it - m_mesh_shader_program_attributes.begin();

	const auto attribute_locations = attribute_location_flags(
		attributes, all_attributes
	);

	m_mesh_shader_program_uniforms.insert(
		m_mesh_shader_program_uniforms.begin() + index, uniforms
	);

	m_mesh_shader_program_attributes.emplace(
		attribute_it, attributes, attribute_locations
	);

	m_mesh_shader_programs.insert(
		m_mesh_shader_programs.begin() + index, shader_program_handle
	);
}


std::optional<shader_program_handle> shader_program_lookup::find(
	shader_program_handle::attribute_support_type attributes,
	shader_program_handle::uniform_support_type uniforms,
	const std::span<const shader_program_variable> all_attributes
) const {

	const auto lower_uniform = std::ranges::lower_bound(
		m_mesh_shader_program_uniforms,
		uniforms
	);

	if (
		lower_uniform == m_mesh_shader_program_uniforms.end() or
		*lower_uniform != uniforms
	) {
		return std::nullopt;
	}

	const auto upper_uniform = std::find_if(
		lower_uniform, m_mesh_shader_program_uniforms.end(),
		[&](const auto& curr_uniforms) {
			return curr_uniforms > uniforms;
		}
	);

	const auto lower_index = lower_uniform - m_mesh_shader_program_uniforms.begin();
	const auto upper_index = upper_uniform - m_mesh_shader_program_uniforms.begin();

	const auto relevant_attributes = std::span(
		m_mesh_shader_program_attributes.begin() + lower_index,
		m_mesh_shader_program_attributes.begin() + upper_index
	);

	const auto upper_attribute = std::upper_bound(
		relevant_attributes.begin(), relevant_attributes.end(),
		attributes,
		[](const auto& attributes, const auto& entry) {
			return attributes < entry.attributes;
		}
	);

	auto entry_it = std::prev(upper_attribute);

	if (
		upper_attribute != relevant_attributes.begin() and
		entry_it->attributes != attributes
	) {

		const auto locations = attribute_location_flags(attributes, all_attributes);

		auto found_match = false;

		while (entry_it != relevant_attributes.begin())
		{
			--entry_it;

			const auto& [ curr_attributes, curr_locations ] = *entry_it;

			// The candidate may not support additional attributes
			if (curr_attributes & ~attributes)
			{
				continue;
			}

			// The candidate may not be missing 'inner' attributes,
			// as this creates location alignment problems.
			const auto higher_neighbour_matched = (curr_locations & locations) >> 1;
			const auto bubbles = higher_neighbour_matched & ~curr_locations;

			if (not bubbles)
			{
				found_match = true;
				break;
			}
		}

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

	const auto shader_program_index = entry_it.base() - m_mesh_shader_program_attributes.begin().base();

	return m_mesh_shader_programs[shader_program_index];
}

shader_program_lookup::attribute_locations_type shader_program_lookup::attribute_location_flags(
	shader_program_handle::attribute_support_type attributes,
	std::span<const shader_program_variable> all_attributes
) {
	auto location_flags = ztu::u32{ 0 };

	auto index = std::size_t{};

	while (attributes)
	{
		if (attributes & 1)
		{
			const auto location = all_attributes[index].info.location;
			location_flags |= attribute_locations_type{ 1 } << location;
		}

		attributes >>= 1;
		++index;
	}

	return location_flags;
}

void shader_program_lookup::print() {
	for (std::size_t i{}; i != m_mesh_shader_program_uniforms.size(); ++i) {
		const auto shader = m_mesh_shader_programs[i];
		const auto uniforms = m_mesh_shader_program_uniforms[i];
		const auto [ attributes, locations ] = m_mesh_shader_program_attributes[i];
		ztu::logger::debug("[%] uniforms: % attributes: % locations: %",
			shader.program_id,
			std::bitset<32>(uniforms),
			std::bitset<32>(attributes),
			std::bitset<32>(locations)
		);
	}
}
}