#include "opengl/data_uploaders/shader_compiler.hpp"

#include <numeric>
#include <bits/ranges_algobase.h>
#include "util/logger.hpp"

void zgl::shader_program_compiler::tokenize_declarations(
	std::string_view source_rest,
	std::vector<std::string_view> tokens,
	std::vector<std::size_t> declaration_token_counts,
	std::span<std::size_t> declaration_type_indices
) {
	tokens.clear();
	declaration_token_counts.clear();
	std::ranges::fill(declaration_type_indices, static_cast<std::size_t>(metadata_declaration_type::invalid));

	constexpr auto pragma_prefix = std::string_view("\n#pragma ");
	constexpr auto title_separator = ':';
	constexpr auto token_separator = ' ';

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

	auto keyword = pragma_prefix;

	while ((offset = source_rest.find(keyword)) != std::string_view::npos)
	{
		const auto current_token_count = tokens.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(title_separator)) == std::string_view::npos)
		{
			continue;
		}

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

		declaration = declaration.substr(offset);

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

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

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

		declaration_token_counts.emplace_back(
			tokens.size() - current_token_count
		);
		source_rest = source_rest.substr(line_end + sizeof('\n'));
		keyword = pragma_prefix.substr(sizeof('\n'));
	}
}

bool zgl::shader_program_compiler::parse_stage_declaration(
	std::span<const std::string_view> tokens,
	shader_program::metadata_type& metadata
) {
	if (tokens.size() != 1)
	{
		ztu::logger::warn("Invalid stage declaration: Expected exactly one token but got %.", tokens.size());
		return false;
	}

	const auto token = tokens.front();

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

	return true;
}

bool zgl::shader_program_compiler::parse_geometry_declaration(
	std::span<const std::string_view> tokens,
	shader_program::metadata_type& metadata
) {
	if (tokens.size() != 1)
	{
		ztu::logger::warn("Invalid geometry declaration: Expected exactly one token but got %.", tokens.size());
		return false;
	}

	const auto token = tokens.front();

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

	return true;
}

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

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

bool zgl::shader_program_compiler::parse_features_declaration(
	std::span<const std::string_view> tokens,
	shader_program::metadata_type& metadata
) {
	switch (metadata.geometry)
	{
	case shader_program::geometry::types::mesh:
		parse_feature_tokens(tokens, mesh_feature_lookup, metadata.features.mesh);
		break;
	case shader_program::geometry::types::point_cloud:
		parse_feature_tokens(tokens, point_cloud_feature_lookup, metadata.features.point_cloud);
		break;
	default:
		ztu::logger::warn("Internal error: Unknown geometry index %.", static_cast<int>(metadata.geometry));
		return false;
	}
	return true;
}

bool zgl::shader_program_compiler::parse_feature_toggles_declaration(
	std::span<const std::string_view> tokens,
	shader_program::metadata_type& metadata
) {
	switch (metadata.geometry)
	{
	case shader_program::geometry::types::mesh:
		parse_feature_tokens(tokens, mesh_feature_lookup, metadata.static_enable.mesh);
		break;
	case shader_program::geometry::types::point_cloud:
		parse_feature_tokens(tokens, point_cloud_feature_lookup, metadata.static_enable.point_cloud);
		break;
	default:
		ztu::logger::warn("Internal error: Unknown geometry index %.", static_cast<int>(metadata.geometry));
		return false;
	}
	return true;
}

std::optional<shader_program::metadata_type> zgl::shader_program_compiler::parse_metadata_from_tokens(
	std::span<const std::string_view> tokens,
	std::span<const std::size_t> declaration_token_counts,
	std::span<const std::size_t> declaration_type_indices
) {
	using enum metadata_declaration_type;
	using namespace std::string_view_literals;

	shader_program::metadata_type data;

	for (const auto [ type, name, parser ] : {
		std::make_tuple(stage, "stage"sv, &parse_stage_declaration),
		std::make_tuple(geometry, "geometry"sv, &parse_geometry_declaration),
		std::make_tuple(features, "features"sv, &parse_features_declaration),
		std::make_tuple(feature_toggles, "feature_toggles"sv, &parse_feature_toggles_declaration)
	}) {
		const auto index = declaration_type_indices[static_cast<std::size_t>(type)];

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

		const auto token_offset = std::accumulate(
			declaration_token_counts.begin(),
			declaration_token_counts.begin() + index,
			std::size_t{}
		);
		const auto token_count = declaration_token_counts[index];

		if (not parser(tokens.subspan(token_offset, token_count), data))
		{
			return std::nullopt;
		}
	}

	return data;
}

auto zgl::shader_program_compiler::find_compatible_shader_source(
	shader_program::metadata_type& requirements
) {

	const auto lower_it = std::ranges::lower_bound(
		shader_lookup,
		requirements,
		shader_program::metadata_type::feature_ignorant_less{},
		&std::pair<shader_program::metadata_type, dynamic_shader_source_store::id_type>::first
	);

	auto generic_requirement_feature_set = requirements.generic_feature_set();
	const auto& required_features = generic_requirement_feature_set.features;

	while (
		lower_it != shader_lookup.end() and
		lower_it->first.geometry == requirements.geometry and
		lower_it->first.stage == requirements.stage
	) {
		const auto& data = lower_it->first;
		const auto& [ features, static_enable, dynamic_enable ] = data.generic_feature_set();

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

		if (missing_features == 0 and fixed_unwanted_features == 0)
		{
			// Tell caller which features need to be toggled before compilation
			// and which features can still be dynamically enabled after compilation.
			generic_requirement_feature_set.static_enable = required_features & static_enable;
			generic_requirement_feature_set.dynamic_enable = dynamic_enable;
			requirements.from_generic_feature_set(generic_requirement_feature_set);

			return lower_it->second;
		}
	}

	return std::nullopt;
}

template<typename T>
void add_required_feature_defines(
	T toggle_flags,
	std::span<const std::string> defines,
	std::vector<const char*>& shader_strings
) {
	auto index = std::size_t{};
	while (toggle_flags != T{})
	{
		if ((toggle_flags & T{ 1 }) != T{})
		{
			shader_strings.push_back(defines[index].c_str());
		}
		toggle_flags >>= 1;
		++index;
	}
}

void zgl::shader_program_compiler::compile_shaders(
	const dynamic_shader_source_store& shader_sources,
	std::span<const shader_program::metadata_type> requirements_list,
	std::vector<shader_handle>& shader_handles
) {

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

	std::vector<const char*> shader_strings;
	shader_strings.reserve(max_shader_strings);


	for (auto requirements : requirements_list)
	{
		auto shader_id = GLuint{};

		const auto source_id = find_compatible_shader_source(requirements);
		if (not source_id)
		{
			continue;
		}

		const auto [ shader_source_it, source_found ] = shader_sources.find(*source_id);
		if (not source_found)
		{
			ztu::logger::warn("Missing shader source with id %.", *source_id);
			continue;
		}

		const auto& shader_source = *shader_source_it;

		if (shader_source.source.empty() and requirements.stage == shader_program::stages::geometry)
		{
			continue;
		}

		shader_strings.clear();

		switch (requirements.geometry) {
		case shader_program::geometry::types::mesh:
				add_required_feature_defines(requirements.feature_set.mesh.static_enable, mesh_feature_defines, shader_strings);
				break;
			case shader_program::geometry::types::point_cloud:
				add_required_feature_defines(requirements.feature_set.point_cloud.static_enable, point_cloud_feature_defines, shader_strings);
				break;
			default:
				std::unreachable();
		}

		shader_strings.push_back(shader_source.source.data());

		shader_id = glCreateShader(static_cast<GLenum>(requirements.stage));

		glShaderSource(shader_id, static_cast<GLsizei>(shader_strings.size()), shader_strings.data(), nullptr);
		glCompileShader(shader_id);

		GLint success;
		glGetShaderiv(shader_id, GL_COMPILE_STATUS, &success);

		if (not success)
		{
			GLint log_length{};
			glGetShaderiv(shader_id, GL_INFO_LOG_LENGTH, &log_length);

			auto log = std::string(log_length, ' ');
			glGetShaderInfoLog(shader_id, log_length, nullptr, log.data());

			ztu::logger::error("Error while compiling shader:\n%", log);

			glDeleteShader(shader_id);
			shader_id = GLuint{};
		}

		shader_handles.emplace_back(shader_id);
	}
}

void zgl::shader_program_compiler::register_shader_sources(
	const dynamic_shader_source_store& shader_sources
) {
	std::vector<std::string_view> tokens;
	std::vector<std::size_t> declaration_token_counts;
	std::array<std::size_t, 4> declaration_type_indices;

	tokens.reserve(32);
	declaration_token_counts.reserve(4);

	for (const auto& [ id, shader_source ] : shader_sources)
	{
		tokenize_declarations(
			shader_source,
			tokens,
			declaration_token_counts,
			declaration_type_indices
		);

		const auto metadata = parse_metadata_from_tokens(
			tokens,
			declaration_token_counts,
			declaration_type_indices
		);

		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(
			shader_lookup,
			*metadata,
			std::less{},
			&std::pair<shader_program::metadata_type, dynamic_shader_source_store::id_type>::first
		);

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

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