#pragma once

#include <unordered_map>
#include <bit>
#include <format>
#include <ranges>
#include <bits/ranges_algo.h>

#include "assets/dynamic_data_stores/dynamic_shader_source_store.hpp"
#include "opengl/shader_program_lookup.hpp"
#include "opengl/handles/shader_handle.hpp"
#include "opengl/handles/shader_program_handle.hpp"
#include "shader_program/metadata_type.hpp"
#include "shader_program/features/mesh_features.hpp"
#include "shader_program/features/point_cloud_features.hpp"
#include "util/string_lookup.hpp"
#include "shader_program/metadata_type.hpp"

namespace zgl
{
class shader_program_compiler
{

public:

	void register_shader_sources(
		const dynamic_shader_source_store& shader_sources
	);

	void compile_shaders(
		const dynamic_shader_source_store& shader_sources,
		std::span<const shader_program::metadata_type> required_capabilities,
		std::vector<shader_handle>& shader_handles
	);


protected:
	void tokenize_declarations(
		std::string_view source,
		std::vector<std::string_view> tokens,
		std::vector<std::size_t> declaration_token_counts,
		std::span<std::size_t> declaration_type_indices
	);

	std::optional<shader_program::metadata_type> 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
	);

	[[nodiscard]] static bool parse_stage_declaration(
		std::span<const std::string_view> tokens,
		shader_program::metadata_type& metadata
	);

	[[nodiscard]] static bool parse_geometry_declaration(
		std::span<const std::string_view> tokens,
		shader_program::metadata_type& metadata
	);

	[[nodiscard]] static bool parse_features_declaration(
		std::span<const std::string_view> tokens,
		shader_program::metadata_type& metadata
	);

	[[nodiscard]] static bool parse_static_enable_declaration(
		std::span<const std::string_view> tokens,
		shader_program::metadata_type& metadata
	);

	[[nodiscard]] static bool parse_dynamic_enable_declaration(
		std::span<const std::string_view> tokens,
		shader_program::metadata_type& metadata
	);

	template<typename T>
	static void parse_feature_tokens(
		std::span<const std::string_view> tokens,
		const ztu::string_lookup<T>& feature_lookup,
		T& features
	);

	std::optional<dynamic_shader_source_store::id_type> zgl::shader_program_compiler::find_compatible_shader_source(
		shader_program::metadata_type& requirements
	);

	template<typename T>
	void add_required_feature_defines(
		T toggle_flags,
		std::span<const std::string> defines,
		std::vector<const char*>& shader_strings
	);

private:
	enum class metadata_declaration_type : std::size_t
	{
		stage = 0,
		geometry = 1,
		features = 2,
		static_enable = 3,
		dynamic_enable = 4,
		invalid = std::numeric_limits<std::size_t>::max()
	};

	inline static auto declaration_lookup = ztu::string_lookup<metadata_declaration_type>{
		{ "STAGE", metadata_declaration_type::stage },
		{ "GEOMETRY", metadata_declaration_type::geometry },
		{ "FEATURES", metadata_declaration_type::features },
		{ "STATIC_ENABLE", metadata_declaration_type::static_enable }
		{ "DYNAMIC_ENABLE", metadata_declaration_type::dynamic_enable }
	};

	inline static auto stage_lookup = ztu::string_lookup<shader_program::stage::types>{
		{ "VERTEX", shader_program::stage::types::vertex },
		{ "GEOMETRY", shader_program::stage::types::geometry },
		{ "FRAGMENT", shader_program::stage::types::fragment },
	};

	inline static auto geometry_lookup = ztu::string_lookup<shader_program::geometry::types>{
		{ "MESH", shader_program::geometry::types::mesh },
		{ "POINT_CLOUD", shader_program::geometry::types::point_cloud }
	};

	inline static auto mesh_feature_lookup = []
	{
		using namespace shader_program::features::mesh;

		auto lookup = ztu::string_lookup<flags>{};
		lookup.reserve(all.size());

		constexpr auto all_flags = std::array{
			flags::face, flags::line, flags::point,
			flags::luminance, flags::color, flags::alpha,
			flags::lighting, flags::texture, flags::uniform_color
		};

		for (const auto& [ define, flag ] : std::ranges::views::zip(defines, all_flags))
		{
			lookup.emplace(std::string(define), flag);
		}

		return lookup;
	}();

	inline static auto point_cloud_feature_lookup = []
	{
		using namespace shader_program::features::point_cloud;

		auto lookup = ztu::string_lookup<flags>{};
		lookup.reserve(all.size());

		constexpr auto all_flags = std::array{
			flags::square, flags::lighting, flags::luminance,
			flags::color, flags::alpha, flags::uniform_color,
			flags::rainbow
		};

		for (const auto& [ define, flag ] : std::ranges::views::zip(defines, all_flags))
		{
			lookup.emplace(std::string(define), flag);
		}

		return lookup;
	}();

	inline static auto mesh_feature_defines = []
	{
		using namespace shader_program::features::mesh;

		auto statements = std::array<std::string, all.size()>{};

		std::ranges::transform(
			defines,
			statements.begin(),
			[](const auto& name) {
				return std::format("#define {}\n", name);
			}
		);

		return statements;
	}();

	inline static auto point_cloud_feature_defines = []
	{
		using namespace shader_program::features::point_cloud;

		auto statements = std::array<std::string, all.size()>{};

		std::ranges::transform(
			defines,
			statements.begin(),
			[](const auto& name) {
				return std::format("#define {}\n", name);
			}
		);

		return statements;
	}();


	std::vector<std::pair<shader_program::metadata_type, dynamic_shader_source_store::id_type>> shader_lookup;
};
}