#include "../../include/viewer/dynamic_shader_program_loading.hpp"
#include "../../include/util/string_lookup.hpp"
#include "../../include/util/logger.hpp"
#include <sstream>

std::size_t viewer::dynamic_shader_program_loading::count_shader_files(
	const std::filesystem::path& path
) {
	using namespace std::string_view_literals;
	namespace fs = std::filesystem;

	auto shader_file_count = std::size_t{ 0 };

	for (const auto& [ asset_type, folder ] : {
		std::make_pair(asset_types::mesh, "mesh"sv),
		std::make_pair(asset_types::point_cloud, "point_cloud"sv)
	}) {
		const auto folder_begin = fs::directory_iterator{ path / folder };
		const auto folder_end = fs::directory_iterator{};

		shader_file_count += std::count_if(
			folder_begin, folder_end,
			[](auto& entry) { return entry.is_regular_file(); }
		);
	}

	return shader_file_count;
}

void viewer::dynamic_shader_program_loading::load_directory(
	asset_loader& loader,
	instance& z3d,
	std::mutex& gl_resource_lock,
	std::mutex& progress_lock,
	std::string& progress_title,
	float& progress_ratio,
	const std::filesystem::path& path
) {

	namespace fs = std::filesystem;
	using namespace std::string_view_literals;

	auto progress_builder = std::stringstream{};

	const auto shader_file_count = static_cast<float>(count_shader_files(path));
	constexpr auto shader_loading_progress = 0.8f;


	auto shader_indices = ztu::string_lookup<ztu::u32>({
		{ "vertex", 0 },
		{ "geometry", 1 },
		{ "fragment", 2 }
	});

	constexpr auto shader_types = std::array<GLenum, 3>{
		GL_VERTEX_SHADER,
		GL_GEOMETRY_SHADER,
		GL_FRAGMENT_SHADER
	};

	auto program_capabilities = std::vector<ztu::u32>{};
	auto programs = std::vector<std::array<zgl::shader_handle, 3>>{};

	auto capability_indices = ztu::string_lookup<ztu::u32>();
	auto capabilities = std::vector<ztu::u32>{};
	capabilities.reserve(8);

	constexpr auto dot_char = '.';
	constexpr auto separator_char = '_';
	constexpr auto optional_char = '?';

	auto curr_shader_count = std::size_t{ 0 };

	for (const auto& [ asset_type, folder ] : {
		std::make_pair(asset_types::mesh, "mesh"sv),
		std::make_pair(asset_types::point_cloud, "point_cloud"sv),
	}) {
		program_capabilities.clear();
		programs.clear();
		capability_indices.clear();

		for (const auto& file : fs::directory_iterator{ path / folder })
		{
			if (not file.is_regular_file())
				continue;

			const auto& file_path = file.path();

			if (file_path.extension() != ".glsl")
				continue;

			const auto filename = file_path.filename().string();

			progress_lock.lock();

			progress_builder.str(std::string{});
			progress_builder << "Loading shader '" <<  filename << "\'...";
			progress_title = progress_builder.str();
			progress_ratio = shader_loading_progress * static_cast<float>(curr_shader_count) / shader_file_count;

			progress_lock.unlock();

			const auto name = std::string_view(filename.begin(), std::ranges::find(filename, dot_char));

			const auto type_str = std::string_view(name.begin(), std::ranges::find(name, separator_char));

			const auto shader_type_index_it = shader_indices.find(type_str);
			if (shader_type_index_it == shader_indices.end()) {
				ztu::logger::warn("Unknown shader type '%'. Skipping shader.", type_str);
				continue;
			}

			const auto shader_type_index = shader_type_index_it->second;
			const auto shader_type = shader_types[shader_type_index];

			auto shader_handle = zgl::shader_handle{};
			if (const auto e = loader.load_shader(shader_type, file_path, shader_handle)) {
				ztu::logger::error(
					"Error while loading shader %: [%] %",
					file_path,
					e.category().name(),
					e.message()
				);
				continue;
			}

			ztu::logger::debug("% -> %", filename, shader_handle.shader_id);

			capabilities.clear();
			capabilities.push_back(0);

			auto specifiers_str = std::string_view(type_str.end() + 1, name.end()); // skip separator_char

			while (not specifiers_str.empty())
			{
				const auto pos = specifiers_str.find(separator_char);
				auto specifier_str = specifiers_str.substr(0, pos);

				if (pos == std::string_view::npos)
				{
					specifiers_str = std::string_view{};
				}
				else
				{
					specifiers_str = specifiers_str.substr(pos + 1); // skip separator_char
				}

				if (specifier_str.empty())
				{
					continue;
				}

				const auto optional = specifier_str.back() == optional_char;
				if (optional)
				{
					specifier_str = specifier_str.substr(0, specifier_str.size() - 1);
				}

				const auto index_it = capability_indices.find(specifier_str);

				auto capability_index = capability_indices.size();
				if (index_it == capability_indices.end())
				{
					capability_indices.emplace(specifier_str, capability_index);
				}
				else
				{
					capability_index = index_it->second;
				}

				const auto capability_flag = ztu::u32{ 1 } << capability_index;

				const auto capability_count = capabilities.size();

				if (optional)
				{
					capabilities.resize(2 * capability_count);
					std::copy_n(capabilities.begin(), capability_count, capabilities.begin() + capability_count);
				}

				for (std::size_t i{}; i != capability_count; ++i)
				{
					capabilities[i] |= capability_flag;
				}
			}

			for (const auto& capability : capabilities)
			{
				const auto program_capability_it = std::ranges::upper_bound(program_capabilities, capability);

				auto program_index = program_capability_it - program_capabilities.begin();

				if (
					program_capability_it == program_capabilities.begin() or
					*std::prev(program_capability_it) != capability
				) {
					program_capabilities.insert(program_capability_it, capability);
					programs.emplace(programs.begin() + program_index);
				}
				else
				{
					--program_index; // The element before the iterator matches.
				}

				programs[program_index][shader_type_index] = shader_handle;
			}

			++curr_shader_count;
		}

		progress_lock.lock();

		progress_title = "Linking programs...";

		progress_lock.unlock();

		// Remove any duplicates shader combinations.
		std::ranges::sort(
			programs,
			[](const auto& lhs, const auto& rhs) {
				return std::lexicographical_compare(
					lhs.begin(), lhs.end(),
					rhs.begin(), rhs.end(),
					[](const auto& a, const auto& b) {
						return a.shader_id < b.shader_id;
					}
				);
			}
		);
		programs.erase(std::ranges::unique(programs).begin(), programs.end());

		ztu::logger::debug("Linking % programs.", programs.size());

		// create shader_program
		for (const auto& [vertex, geometry, fragment] : programs)
		{
			if (vertex.shader_id == 0 or fragment.shader_id == 0)
			{
				ztu::logger::warn(
					"Skipping program as the combination is unlikely to be used (vertex: % geometry: % fragment: %).",
					vertex.shader_id, geometry.shader_id, fragment.shader_id
				);
				continue;
			}

			auto program_handle = zgl::shader_program_handle{};
			if (const auto e = loader.build_shader_program(vertex, geometry, fragment, program_handle))
			{
				ztu::logger::error(
					"Error occurred while linking shader program: [%] %",
					e.category().name(),
					e.message()
				);
				continue;
			}

			ztu::logger::debug(
				"Linked (vertex: % geometry: % fragment: %) -> %",
				vertex.shader_id, geometry.shader_id, fragment.shader_id,
				program_handle.program_id
			);


			gl_resource_lock.lock();
			z3d.add_shader_program(asset_type, program_handle);
			gl_resource_lock.unlock();
		}

		gl_resource_lock.lock();
		z3d.m_mesh_shader_program_lookup.print();
		gl_resource_lock.unlock();

		loader.unload_shader_data();
	}
}