#include "opengl/data_managers/shader_source_manager.hpp"

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

#include "opengl/shading/shader_metadata_language.hpp"

static auto mesh_feature_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_feature_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_feature_set_comparator
{
	using type = zgl::shader_features_set<zgl::shading::features::generic::type>;

	bool operator()(const type& a, const type& b) const noexcept
	{
		static constexpr auto more_features = std::popcount<zgl::shading::features::generic::type>;

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

struct prioritized_metadata_comparator
{
	using type = zgl::shader_source_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;
		}

		const auto features_a = a.generic_feature_set();
		const auto features_b = b.generic_feature_set();

		return feature_set_comparator(features_a, features_b);
	}

private:
	prioritized_feature_set_comparator feature_set_comparator{};
};



void zgl::shader_source_manager::process(
	const assets::data_stores& stores
) {
	namespace language = shading::shader_metadata_language;

	for (const auto& [ id, shader_source ] : stores.shader_sources)
	{
		m_value_token_buffer.clear();
		m_declaration_token_count_buffer.clear();
		std::ranges::fill(
			m_declaration_type_index_buffer,
			static_cast<std::size_t>(language::declaration_type::invalid)
		);

		tokenize_declarations(shader_source);

		const auto metadata = parse_metadata_from_tokens();

		if (not metadata)
		{
			ztu::logger::warn("Ignoring shader % as it contains malformed metadata.", id);
			continue;
		}

		// 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,
			*metadata,
			prioritized_metadata_comparator{},
			&std::pair<shader_source_metadata, dynamic_shader_source_store::id_type>::first
		);

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

		m_shader_source_lookup.emplace(it, *metadata, id);
	}
}

void zgl::shader_source_manager::get_shader_sources(
	const assets::shader_source_store& shader_sources,
	std::span<const shading::shader_source_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_feature_defines.size(),
		point_cloud_feature_defines.size()
	) + 1;

	shader_strings.reserve(max_shader_strings);

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

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

			assets::shader_source_store::id_type source_id{};
			shading::features::generic::type to_be_enabled{};

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

				const auto missing_features = req.features & ~features;
				const auto unwanted_features = ~req.features & features;
				const auto fixed_unwanted_features = unwanted_features & ~static_enable & ~dynamic_enable;

				if (missing_features == 0 and fixed_unwanted_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;
				}

				++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,
						to_be_enabled,
						res.string_count,
						shader_strings
					);
				}
			}

			return res;
		}
	);
}

void zgl::shader_source_manager::get_define_strings(
	const shading::model_geometry::types geometry,
	shading::features::generic::type features,
	shading::features::generic::type& feature_count,
	std::vector<const char*>& defines
) {
	std::span<const char*> all_defines;

	switch (geometry)
	{
		case shading::model_geometry::types::mesh:
			all_defines = mesh_feature_defines;
			break;
		case shading::model_geometry::types::point_cloud:
			all_defines = point_cloud_feature_defines;
			break;
		default:
			std::unreachable();
	}

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


void zgl::shader_source_manager::tokenize_declarations(
	std::string_view source_rest
) {
	namespace language = shading::shader_metadata_language;

	auto offset = std::string_view::size_type{};

	auto keyword = language::declaration_prefix;

	while ((offset = source_rest.find(keyword)) != std::string_view::npos)
	{
		const auto current_token_count = m_value_token_buffer.size();

		auto line_end = source_rest.find('\n', offset);
		if (line_end == std::string_view::npos)
		{
			line_end = source_rest.length();
		}

		auto declaration = source_rest.substr(offset, line_end - offset);

		if ((offset = declaration.find(language::title_separator)) == std::string_view::npos)
		{
			continue;
		}

		const auto title = declaration.substr(0, offset);
		if (const auto it = language::declaration_lookup.find(title); it != language::declaration_lookup.end())
		{
			const auto declaration_type = static_cast<std::size_t>(it->second);
			m_declaration_type_index_buffer[declaration_type] = m_declaration_token_count_buffer.size();
		}
		else
		{
			continue;
		}

		declaration = declaration.substr(offset);

		if (not declaration.empty() and declaration.front() == language::value_separator)
		{
			declaration = declaration.substr(sizeof(language::value_separator), declaration.length());
		}

		while ((offset = declaration.find(language::value_separator)) != std::string_view::npos)
		{
			m_value_token_buffer.emplace_back(declaration.substr(0, offset));
			declaration = declaration.substr(offset + sizeof(language::value_separator));
		}

		if (not declaration.empty())
		{
			m_value_token_buffer.emplace_back(declaration);
		}

		m_declaration_token_count_buffer.emplace_back(
			m_value_token_buffer.size() - current_token_count
		);
		source_rest = source_rest.substr(line_end + sizeof('\n'));
		keyword = language::declaration_prefix.substr(sizeof('\n'));
	}
}

bool zgl::shader_source_manager::parse_stage_declaration(
	std::span<const std::string_view> values,
	shader_source_metadata& metadata
) {
	namespace language = shading::shader_metadata_language;

	if (values.size() != 1)
	{
		ztu::logger::warn("Invalid stage declaration: Expected exactly one token but got %.", values.size());
		return false;
	}

	const auto value = values.front();

	if (const auto it = language::stage_lookup.find(value); it != language::stage_lookup.end())
	{
		metadata.stage = it->second;
	}
	else
	{
		ztu::logger::warn("Invalid stage declaration: Unknown stage %.", value);
		return false;
	}

	return true;
}

bool zgl::shader_source_manager::parse_geometry_declaration(
	std::span<const std::string_view> values,
	shader_source_metadata& metadata
) {
	namespace language = shading::shader_metadata_language;

	if (values.size() != 1)
	{
		ztu::logger::warn("Invalid geometry declaration: Expected exactly one token but got %.", values.size());
		return false;
	}

	const auto value = values.front();

	if (const auto it = language::geometry_lookup.find(value); it != language::geometry_lookup.end())
	{
		metadata.geometry = it->second;
	}
	else
	{
		ztu::logger::warn("Invalid geometry declaration: Unknown geometry %.", value);
		return false;
	}

	return true;
}

template<typename T>
void zgl::shader_source_manager::parse_feature_tokens(
	std::span<const std::string_view> values,
	const ztu::string_lookup<T>& feature_lookup,
	T& features
) {
	features = {};

	for (const auto value : values)
	{
		if (const auto it = feature_lookup.find(value); it != feature_lookup.end())
		{
			features |= it->second;
		}
		else
		{
			ztu::logger::warn("Ignoring unknown feature token %.", value);
		}
	}
}

bool zgl::shader_source_manager::parse_features_declaration(
	std::span<const std::string_view> values,
	shader_source_metadata& metadata
) {
	namespace language = shading::shader_metadata_language;

	switch (metadata.geometry)
	{
	case shading::model_geometry::types::mesh:
		parse_feature_tokens(values, language::mesh_feature_lookup, metadata.feature_set.mesh.features);
		break;
	case shading::model_geometry::types::point_cloud:
		parse_feature_tokens(values, language::point_cloud_feature_lookup, metadata.feature_set.point_cloud.features);
		break;
	default:
		ztu::logger::warn("Internal error: Unknown geometry index %.", static_cast<int>(metadata.geometry));
		return false;
	}
	return true;
}

bool zgl::shader_source_manager::parse_static_enable_declaration(
	std::span<const std::string_view> values,
	shader_source_metadata& metadata
) {
	namespace language = shading::shader_metadata_language;

	switch (metadata.geometry)
	{
	case shading::model_geometry::types::mesh:
		parse_feature_tokens(values, language::mesh_feature_lookup, metadata.feature_set.mesh.static_enable);
		break;
	case shading::model_geometry::types::point_cloud:
		parse_feature_tokens(values, language::point_cloud_feature_lookup, metadata.feature_set.point_cloud.static_enable);
		break;
	default:
		ztu::logger::warn("Internal error: Unknown geometry index %.", static_cast<int>(metadata.geometry));
		return false;
	}
	return true;
}

bool zgl::shader_source_manager::parse_dynamic_enable_declaration(
	std::span<const std::string_view> values,
	shader_source_metadata& metadata
) {
	namespace language = shading::shader_metadata_language;

	switch (metadata.geometry)
	{
	case shading::model_geometry::types::mesh:
		parse_feature_tokens(values, language::mesh_feature_lookup, metadata.feature_set.mesh.dynamic_enable);
		break;
	case shading::model_geometry::types::point_cloud:
		parse_feature_tokens(values, language::point_cloud_feature_lookup, metadata.feature_set.point_cloud.dynamic_enable);
		break;
	default:
		ztu::logger::warn("Internal error: Unknown geometry index %.", static_cast<int>(metadata.geometry));
		return false;
	}
	return true;
}

std::optional<zgl::shader_source_metadata> zgl::shader_source_manager::parse_metadata_from_tokens()
{
	namespace language = shading::shader_metadata_language;
	using namespace std::string_view_literals;

	shader_source_metadata data;

	for (const auto [ type, name, parser ] :
		{
			std::make_tuple(language::declaration_type::stage, "stage"sv, &parse_stage_declaration),
			std::make_tuple(language::declaration_type::geometry, "geometry"sv, &parse_geometry_declaration),
			std::make_tuple(language::declaration_type::features, "features"sv, &parse_features_declaration),
			std::make_tuple(language::declaration_type::static_enable, "static_enable"sv, &parse_static_enable_declaration),
			std::make_tuple(language::declaration_type::dynamic_enable, "dynamic_enable"sv, &parse_dynamic_enable_declaration)
		}
	) {
		const auto index = m_declaration_type_index_buffer[static_cast<std::size_t>(type)];

		if (index == static_cast<std::size_t>(language::declaration_type::invalid))
		{
			ztu::logger::warn("Shader metadata error: Missing % declaration.", name);
			return std::nullopt;
		}

		const auto value_token_offset = std::accumulate(
			m_declaration_token_count_buffer.begin(),
			m_declaration_token_count_buffer.begin() + index,
			std::size_t{}
		);
		const auto value_token_count = m_declaration_token_count_buffer[index];

		if (not parser(std::span(m_value_token_buffer).subspan(value_token_offset, value_token_count), data))
		{
			return std::nullopt;
		}
	}

	return data;
}