#include "opengl/data_managers/shader_source_manager.hpp"

#include "util/logger.hpp"
#include <algorithm>
#include <cassert>
#include <numeric>
#include <utility>

#include "opengl/shading/shader_metadata_language.hpp"

static auto mesh_component_defines = std::array{
	"#define FACE\n",
	"#define LINE\n",
	"#define POINT\n",
	"#define V_L\n",
	"#define V_RGB\n",
	"#define V_A\n",
	"#define LIGHTING\n",
	"#define TEXTURE\n",
	"#define U_RGBA\n",
};

static auto point_cloud_component_defines = std::array{
	"#define SQUARE\n",
	"#define LIGHTING\n",
	"#define V_L\n",
	"#define V_RGB\n",
	"#define V_A\n",
	"#define U_RGBA\n",
	"#define RAINBOW\n"
};

struct prioritized_metadata_comparator
{
	using type = zgl::shader_source_metadata;

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

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

		static constexpr auto more_components = std::popcount<assets::shader_components::flags>;

		// Sort by dynamic components first to make sure when compatible components are found
		// the compiled shader will have maximum dynamic compatibility.

		return std::ranges::lexicographical_compare(
			std::array{ a.dynamic_enable, a.components, a.static_enable },
			std::array{ b.dynamic_enable, b.components, b.static_enable },
			std::greater{},
			more_components,
			more_components
		);
	}
};


void zgl::shader_source_manager::process(const store_type& shader_sources)
{
	namespace language = shading::shader_metadata_language;

	for (const auto& [ id, shader_source ] : shader_sources)
	{
		// Sorted insert should be faster than std::sort and std::unique
		// for small numbers of elements and high numbers of duplicates.
		const auto it = std::ranges::upper_bound(
			m_shader_source_lookup,
			shader_source.meta,
			prioritized_metadata_comparator{},
			&entry_type::first
		);

		if (it != m_shader_source_lookup.end() and it->first == shader_source.meta)
		{
			continue;
		}

		m_shader_source_lookup.emplace(it, shader_source.meta, id);
	}
}

void zgl::shader_source_manager::fetch(
	const assets::shader_source_store& shader_sources,
	std::span<const shading::shader_requirements> requirements,
	std::span<preprocessed_shader_source_metadata> metadata,
	std::vector<const char*>& shader_strings
) {

	assert(requirements.size() == metadata.size());

	static constexpr auto max_shader_strings = std::max(
		mesh_component_defines.size(),
		point_cloud_component_defines.size()
	) + 1;

	shader_strings.reserve(max_shader_strings);

	std::ranges::transform(
		requirements,
		metadata.begin(),
		[&](const shading::shader_requirements& req)
		{
			auto res = preprocessed_shader_source_metadata{};

			auto source_it = std::ranges::lower_bound(
				m_shader_source_lookup,
				std::pair{ req.geometry_type, req.stage },
				std::greater{},
				[](const entry_type& entry)
				{
					const auto& meta = entry.first;
					return std::pair{ meta.geometry_type, meta.stage };
				}
			);

			assets::shader_source_store::id_type source_id{};
			assets::shader_components::flags to_be_enabled{};

			while (
				source_it != m_shader_source_lookup.end() and
				source_it->first.geometry_type == req.geometry_type and
				source_it->first.stage == req.stage
			) {
				const auto& [ meta, id ] = *source_it;

				const auto missing_components = req.components & ~meta.components;
				const auto unwanted_components = ~req.components & meta.components;
				const auto fixed_unwanted_components = unwanted_components & ~meta.static_enable & ~meta.dynamic_enable;

				if (missing_components == 0 and fixed_unwanted_components == 0)
				{
					to_be_enabled = req.components & meta.static_enable;
					source_id = id;

					res.static_enabled = meta.components & ~meta.dynamic_enable & ~unwanted_components;
					res.dynamic_enable = meta.dynamic_enable;

					break;
				}

				++source_it;
			}

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

			return res;
		}
	);
}

void zgl::shader_source_manager::get_define_strings(
	const assets::model_geometry::types geometry_type,
	assets::shader_components::flags components,
	assets::shader_components::flags& component_count,
	std::vector<const char*>& defines
) {
	std::span<const char*> all_defines;

	switch (geometry_type)
	{
		case shading::model_geometry::types::mesh:
			all_defines = mesh_component_defines;
			break;
		case shading::model_geometry::types::point_cloud:
			all_defines = point_cloud_component_defines;
			break;
		default:
			std::unreachable();
	}

	auto index = std::size_t{};
	while (components != 0)
	{
		if ((components & 1) != 0)
		{
			defines.push_back(all_defines[index]);
			++component_count;
		}
		components >>= 1;
		++index;
	}
}