#include "assets/file_parsers/glsl_parser.hpp"

#include <fstream>
#include <numeric>

#include "assets/components/mesh_shader_components.hpp"
#include "assets/components/point_cloud_shader_components.hpp"
#include "util/logger.hpp"
#include <mutex>
#include <execution>

#include "util/string_lookup.hpp"

namespace assets::language
{
	enum class 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()
	};

	using stage = shader_components::stage;

	inline constexpr auto declaration_prefix = std::string_view("\n#pragma ");
	inline constexpr auto title_separator = ':';
	inline constexpr auto value_separator = ' ';

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

	inline auto stage_lookup = ztu::string_lookup<shader_components::stage>{
		{ "VERTEX", stage::vertex },
		{ "TESSELATION_CONTROL", stage::tesselation_control },
		{ "TESSELATION_EVALUATION", stage::tesselation_evaluation },
		{ "GEOMETRY", stage::geometry },
		{ "FRAGMENT", stage::fragment },
	};

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

	inline auto mesh_feature_lookup = ztu::string_lookup<shader_components::flags>{
		{ "FACE", static_cast<shader_components::flags>(mesh_shader_components::flags::face) },
		{ "LINE", static_cast<shader_components::flags>(mesh_shader_components::flags::line) },
		{ "POINT", static_cast<shader_components::flags>(mesh_shader_components::flags::point) },
		{ "V_L", static_cast<shader_components::flags>(mesh_shader_components::flags::luminance) },
		{ "V_RGB", static_cast<shader_components::flags>(mesh_shader_components::flags::color) },
		{ "V_R", static_cast<shader_components::flags>(mesh_shader_components::flags::alpha) },
		{ "COLOR_TEXTURE", static_cast<shader_components::flags>(mesh_shader_components::flags::color_texture) },
		{ "U_LIGHTING", static_cast<shader_components::flags>(mesh_shader_components::flags::uniform_lighting) },
		{ "T_LIGHTING", static_cast<shader_components::flags>(mesh_shader_components::flags::textured_lighting) },
		{ "U_RGBA", static_cast<shader_components::flags>(mesh_shader_components::flags::uniform_color) }
	};

	inline auto point_cloud_feature_lookup = ztu::string_lookup<shader_components::flags>{
		{ "SQUARE", static_cast<shader_components::flags>(point_cloud_shader_components::flags::square) },
		{ "LIGHTING", static_cast<shader_components::flags>(point_cloud_shader_components::flags::lighting) },
		{ "V_L", static_cast<shader_components::flags>(point_cloud_shader_components::flags::luminance) },
		{ "V_RGB", static_cast<shader_components::flags>(point_cloud_shader_components::flags::color) },
		{ "V_A", static_cast<shader_components::flags>(point_cloud_shader_components::flags::alpha) },
		{ "U_RGBA", static_cast<shader_components::flags>(point_cloud_shader_components::flags::uniform_color) },
		{ "RAINBOW", static_cast<shader_components::flags>(point_cloud_shader_components::flags::rainbow) }
	};
}


std::error_code assets::glsl_parser::prefetch(
	path_id_lookups&
) {
	// Nothing to prefetch
	return {};
}

std::error_code assets::glsl_parser::load(
	const path_id_lookups& lookups,
	shader_source_store& store,
	bool
) {
	m_path_buffer.clear();
	lookups.shader_sources.by_extension(".glsl", m_path_buffer);

	std::for_each(
		std::execution::parallel_unsequenced_policy{},
		m_path_buffer.begin(),
		m_path_buffer.end(),
		parser_context{ store }
	);

	return {};
}


assets::glsl_parser::parser_context::parser_context(
	store_type& m_store
) : m_store{ &m_store }
{
	m_buffer.source.reserve(4096);
	m_value_buffer.reserve(64);
	m_declaration_value_count_buffer.reserve(8);
}

void assets::glsl_parser::parser_context::reset()
{
	m_buffer.clear();
	m_value_buffer.clear();
	m_declaration_value_count_buffer.clear();
	std::ranges::fill(
		m_declaration_type_index_buffer,
		static_cast<std::size_t>(language::declaration_type::invalid)
	);
}

void assets::glsl_parser::parser_context::operator()(lookup_type::const_pointer entry) noexcept
{
	const auto& [ filename, id ] = *entry;

	if (const auto e = read_file(filename, m_buffer.source))
	{
		ztu::logger::warn("Could not load shader source %: %.", filename, e.message());
		return;
	}

	tokenize_declarations();

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

	remove_metadata_declarations();

	m_store->insert(id, m_buffer);
}


void assets::glsl_parser::parser_context::tokenize_declarations()
{
	auto source_rest = std::string_view{ m_buffer.source.data(), m_buffer.source.size() };
	auto offset = std::string_view::size_type{};

	while ((offset = source_rest.find(language::declaration_prefix)) != std::string_view::npos)
	{
		const auto current_token_count = m_value_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_value_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_buffer.emplace_back(declaration.substr(0, offset));
			declaration = declaration.substr(offset + sizeof(language::value_separator));
		}

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

		m_declaration_value_count_buffer.emplace_back(
			m_value_buffer.size() - current_token_count
		);

		// Preserve line break
		source_rest = source_rest.substr(line_end);
	}
}

bool assets::glsl_parser::parser_context::parse_metadata_from_tokens()
{
	using namespace std::string_view_literals;

	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_components_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 false;
		}

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

		if (not parser(std::span(m_value_buffer).subspan(value_token_offset, value_token_count), m_buffer.meta))
		{
			return false;
		}
	}

	return true;
}

void assets::glsl_parser::parser_context::remove_metadata_declarations()
{

	if (m_declaration_value_count_buffer.empty())
	{
		return;
	}

	// 1. For every declaration expand the first value view to the size of the line.
	// 2. Delete all other value views as they are not needed anymore.
	// 3. Delete all declaration lines by overwriting them with the following source code.
	// 4. Resize the source vector to the new character count.

	const auto source_view = std::string_view{ m_buffer.source.data(), m_buffer.source.size() };
	const auto value_span = std::span(m_value_buffer);

	auto offset = std::size_t{};
	for (auto [ index, count ] : std::ranges::views::enumerate(m_declaration_value_count_buffer))
	{
		const auto values = value_span.subspan(offset, count);
		const auto value = values.front();

		std::size_t begin_pos = value.data() - source_view.data();
		std::size_t end_pos = begin_pos + value.size();

		begin_pos = source_view.rfind('\n', begin_pos);
		if (begin_pos == std::string_view::npos)
		{
			begin_pos = 0;
		}

		end_pos = source_view.find('\n', end_pos);
		if (end_pos == std::string_view::npos)
		{
			end_pos = source_view.length();
		}
		else
		{
			end_pos += sizeof('\n');
		}

		m_value_buffer[index] = source_view.substr(begin_pos, end_pos - begin_pos);

		offset += count;
	}

	const auto lines = std::span(m_value_buffer.data(), m_declaration_value_count_buffer.size());

	const auto& first_line = lines.front();
	auto source_end = first_line.begin();
	auto line_end = first_line.end();

	for (const auto& line : lines.subspan(1))
	{
		std::copy(line_end, line.begin(), source_end);
		source_end += line.size();
		line_end = line.end();
	}

	const auto source_rest_size = m_buffer.source.end().base() - line_end;
	std::copy_n(line_end, source_rest_size, source_end);

	source_end += source_rest_size;

	m_buffer.source.resize(source_end - m_buffer.source.data());
}



std::error_code assets::glsl_parser::read_file(
	const std::filesystem::path& filename,
	std::vector<char>& source
) {
	auto file = std::ifstream{ filename };

	if (not file.is_open())
	{
		return std::make_error_code(std::errc::no_such_file_or_directory);
	}

	file.seekg(0, std::ios::end);
	const std::streampos size = file.tellg();

	if (size == 0 or size == std::numeric_limits<std::streampos>::max())
	{
		return std::make_error_code(std::errc::invalid_seek);
	}

	file.seekg(0, std::ios::beg);

	source.clear();
	source.reserve(sizeof('\n') + size);
	source.push_back('\n');

	std::copy(
		std::istreambuf_iterator<char>(file),
		std::istreambuf_iterator<char>(),
		std::back_inserter(source)
	);

	file.close();

	return {};
}

bool assets::glsl_parser::parse_geometry_declaration(
	const std::span<const std::string_view> values,
	shader_source_data::metadata& meta
) {

	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())
	{
		meta.geometry_type = it->second;
	}
	else
	{
		ztu::logger::warn("Invalid geometry declaration: Unknown geometry %.", value);
		return false;
	}

	return true;
}

bool assets::glsl_parser::parse_stage_declaration(
	const std::span<const std::string_view> values,
	shader_source_data::metadata& meta
) {

	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())
	{
		meta.stage = it->second;
	}
	else
	{
		ztu::logger::warn("Invalid stage declaration: Unknown stage %.", value);
		return false;
	}

	return true;
}

bool assets::glsl_parser::parse_components_declaration(
	const std::span<const std::string_view> values,
	shader_source_data::metadata& meta
) {
	return parse_component_tokens(values, meta.geometry_type, meta.components);
}

bool assets::glsl_parser::parse_static_enable_declaration(
	const std::span<const std::string_view> values,
	shader_source_data::metadata& meta
) {
	return parse_component_tokens(values, meta.geometry_type, meta.static_enable);
}

bool assets::glsl_parser::parse_dynamic_enable_declaration(
	const std::span<const std::string_view> values,
	shader_source_data::metadata& meta
) {
	return parse_component_tokens(values, meta.geometry_type, meta.dynamic_enable);
}

bool assets::glsl_parser::parse_component_tokens(
	const std::span<const std::string_view> values,
	const model_geometry::types geometry_type,
	shader_components::flags& components
) {
	components = {};

	const ztu::string_lookup<shader_components::flags>* component_lookup{};

	switch (geometry_type)
	{
	case model_geometry::types::mesh:
		component_lookup = &language::mesh_feature_lookup;
		break;
	case model_geometry::types::point_cloud:
		component_lookup = &language::point_cloud_feature_lookup;
		break;
	default:
		ztu::logger::warn("Internal error: Unknown geometry index %.", static_cast<int>(geometry_type));
		return false;
	}


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

	return true;
}