#include "assets/file_parsers/obj_parser.hpp"

#include <charconv>
#include <fstream>
#include <array>

#include "assets/components/mesh_vertex_components.hpp"

#include "util/logger.hpp"
#include "util/for_each.hpp"
#include "util/line_parser.hpp"
#include <execution>

namespace assets::obj_parser_error
{
struct category : std::error_category
{
	[[nodiscard]] const char* name() const noexcept override {
		return "obj_loader";
	}

	[[nodiscard]] std::string message(int ev) const override
	{
		switch (static_cast<codes>(ev))
		{
			using enum codes;
		case cannot_open_file:
			return "Cannot open given obj file.";
		case malformed_vertex:
			return "File contains malformed 'v' statement.";
		case malformed_texture_coordinate:
			return "File contains malformed 'vt' statement.";
		case malformed_normal:
			return "File contains malformed 'vn' statement.";
		case malformed_face:
			return "File contains malformed 'f' statement.";
		case face_index_out_of_range:
			return "Face index out of range.";
		case unknown_line_begin:
			return "Unknown obj line begin.";
		case use_material_without_material_library:
			return "'usemtl' statement before material library loaded.";
		case unknown_material_name:
			return "No matching material name found in material library.";
		default:
			using namespace std::string_literals;
			return "unrecognized error ("s + std::to_string(ev) + ")";
		}
	}
};

} // namespace mesh_loader_error

inline std::error_category& obj_loader_error_category()
{
	static assets::obj_parser_error::category category;
	return category;
}

namespace assets::obj_parser_error
{
inline std::error_code make_error_code(codes e)
{
	return { static_cast<int>(e), obj_loader_error_category() };
}
} // namespace mesh_loader_error


std::error_code assets::obj_parser::prefetch(
	path_id_lookups& lookups
) {
	m_path_buffer.clear();
	lookups.meshes.by_extension(".obj", m_path_buffer);

	std::for_each(
		std::execution::parallel_unsequenced_policy{},
		m_path_buffer.begin(),
		m_path_buffer.end(),
		prefetcher_context{
			lookups
		}
	);

	return {};
}

std::error_code assets::obj_parser::load(
	path_id_lookups& lookups,
	data_stores& stores,
	bool pedantic
) {
	m_path_buffer.clear();
	lookups.meshes.by_extension(".obj", m_path_buffer);

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

	return {};
}

template<int L, typename T>
std::errc parse_numeric_vector(std::string_view param, z3d::vec<L, T>& vec)
{
	auto it = param.begin();
	const auto end = param.end();

	for (int i{}; i != L; ++i)
	{
		if (it >= end)
		{
			return std::errc::invalid_argument;
		}

		const auto [ptr, ec] = std::from_chars(it, end, vec[i]);

		if (ec != std::errc{})
		{
			return ec;
		}

		it = ptr + 1; // Skip space in between components.
	}

	return {};
};


assets::obj_parser::parser_context::parser_context(
	path_id_lookups& m_id_lookups,
	data_stores& m_stores
) :
	m_id_lookups{ &m_id_lookups },
	m_stores{ &m_stores }
{
	constexpr auto expected_vertex_count = 8192;
	m_mesh.positions().reserve(expected_vertex_count);
	m_mesh.normals().reserve(expected_vertex_count);
	m_mesh.colors().reserve(expected_vertex_count);
	m_mesh.reflectances().reserve(expected_vertex_count);
	m_mesh.tex_coords().reserve(expected_vertex_count);
	m_mesh.triangles().reserve(2 * expected_vertex_count);

	m_position_buffer.reserve(expected_vertex_count);
	m_normal_buffer.reserve(expected_vertex_count);
	m_tex_coord_buffer.reserve(expected_vertex_count);
}

void assets::obj_parser::parser_context::reset()
{
	m_mesh.clear();
	m_position_buffer.clear();
	m_normal_buffer.clear();
	m_tex_coord_buffer.clear();
	m_vertex_comp_indices_to_vertex_index.clear();
}


void assets::obj_parser::parser_context::operator()(lookup_type::const_pointer entry) noexcept
{
	using obj_parser_error::codes;
	using obj_parser_error::make_error_code;
	namespace fs = std::filesystem;

	reset();

	const auto& [ filename, id ] = *entry;

	auto path_buffer = fs::path{};
	const auto base_dir = fs::canonical(fs::path(filename).parent_path());

	const auto push_mesh = [&](const bool clear_read_buffer = false)
	{
		if (not m_mesh.triangles().empty())
		{
			ztu::logger::debug("parsed % positions.", m_position_buffer.size());
			ztu::logger::debug("parsed % normals.", m_normal_buffer.size());
			ztu::logger::debug("parsed % tex_coords.", m_tex_coord_buffer.size());

			ztu::logger::debug("stored % positions.", m_mesh.positions().size());
			ztu::logger::debug("stored % normals.", m_mesh.normals().size());
			ztu::logger::debug("stored % tex_coords.", m_mesh.tex_coords().size());
			ztu::logger::debug("stored % triangles.", m_mesh.triangles().size());

			m_stores->meshes.insert(id, m_mesh);
		}

		if (clear_read_buffer)
		{
			m_position_buffer.clear();
			m_normal_buffer.clear();
			m_tex_coord_buffer.clear();
		}

		m_mesh.clear();
		m_vertex_comp_indices_to_vertex_index.clear();
	};

	const material_library_data* curr_material_library{};


	auto in = std::ifstream{ filename };
	if (not in.is_open())
	{
		ztu::logger::warn("Cannot open obj file %.", filename);
		return;
	}

	const auto ec = ztu::parse_lines<codes>(
		in,
		false,
		make_line_parser("v ", ztu::is_repeating, [&](const auto& param)
		{
			mesh_vertex_components::position position;
			if (parse_numeric_vector(param, position) != std::errc{}) [[unlikely]]
			{
				return codes::malformed_vertex;
			}

			m_position_buffer.push_back(position);

			return codes::ok;
		}),
		make_line_parser("vt ", ztu::is_repeating, [&](const auto& param)
		{
			mesh_vertex_components::tex_coord coord;
			if (parse_numeric_vector(param, coord) != std::errc{}) [[unlikely]]
			{
				return codes::malformed_texture_coordinate;
			}

			m_tex_coord_buffer.push_back(coord);

			return codes::ok;
		}),
		make_line_parser("vn ", ztu::is_repeating, [&](const auto& param)
		{
			mesh_vertex_components::normal normal;
			if (parse_numeric_vector(param, normal) != std::errc{}) [[unlikely]]
			{
				return codes::malformed_normal;
			}

			m_normal_buffer.push_back(normal);

			return codes::ok;
		}),
		make_line_parser("o ", ztu::is_not_repeating, [&](const auto&)
		{
			push_mesh(); // Name is currently ignored

			return codes::ok;
		}),
		make_line_parser("f ", ztu::is_repeating, [&](const auto& param)
		{
			const auto begin = param.begin().base();
			const auto end = param.end().base();

			z3d::vertex_index first_index{}, prev_index{};

			auto vertex_count = std::size_t{};

			for (auto it = begin; it <= end; ++it)
			{
				auto vertex = component_indices{};

				for (auto& component_index : vertex)
				{
					if (it != end and *it == '/')
					{
						++it;
						continue;
					}

					const auto [ptr, ec] = std::from_chars(it, end, component_index);
					if (ec != std::errc()) [[unlikely]]
					{
						// Discard whole face if one index is malformed.
						return codes::malformed_face;
					}

					it = ptr;

					if (it == end or *it != '/')
					{
						break;
					}

					++it;
				}

				++vertex_count;

				if (it != end and *it != ' ') [[unlikely]]
				{
					return codes::malformed_face;
				}

				const auto curr_index = find_or_push_vertex(vertex);

				if (vertex_count >= 3)
				{
					m_mesh.triangles().emplace_back() = {
						first_index,
						prev_index,
						curr_index
					};
				}
				else if (vertex_count == 1)
				{
					first_index = curr_index;
				}

				prev_index = curr_index;
		   }

			return codes::ok;
		}),
		make_line_parser("g ", ztu::is_not_repeating, [&](const auto& param)
		{
			push_mesh(false);  // Name is currently ignored

			return codes::ok;
		}),
		make_line_parser("usemtl ", ztu::is_not_repeating, [&](const auto& param)
		{
			push_mesh(false);

			if (not curr_material_library) [[unlikely]]
			{
				return codes::use_material_without_material_library;
			}

			const auto material_id_it = curr_material_library->find(param);

			if (material_id_it == curr_material_library->end()) [[unlikely]]
			{
				return codes::unknown_material_name;
			}

			m_mesh.material() = material_id_it->second;

			return codes::ok;
		}),
		make_line_parser("mtllib ", ztu::is_not_repeating, [&](const auto& param)
		{
			path_buffer.assign(param);

			if (path_buffer.is_relative())
			{
				path_buffer = base_dir;
				path_buffer /= param;
			}

			const auto material_library_id_it = m_id_lookups->material_libraries.find(path_buffer);

			if (material_library_id_it != m_id_lookups->material_libraries.end()) [[likely]]
			{
				const auto material_library_id = material_library_id_it->second;

				const auto [ it, found ] = m_stores->material_libraries.find(material_library_id);
				if (found)
				{
					// TODO implement proper dereference operator
					curr_material_library = &((*it).second);
				}
				else
				{
					ztu::logger::warn(
						"Material with name '%' was not loaded and can therefor not be used.",
						param
					);
				}
			}
			else [[unlikely]]
			{
				ztu::logger::warn(
					"No material with name '%' found in the material library.",
					param
				);
				curr_material_library = nullptr;
			}

			return codes::ok;
		})
	);

	if (ec != codes::ok)
	{
		const auto e = make_error_code(ec);
		ztu::logger::error("Error while parsing obj file %: %", filename, e.message());
	}

	push_mesh();
}


z3d::vertex_index assets::obj_parser::parser_context::find_or_push_vertex(const component_indices& vertex_comp_indices)
{
	// If the triangle indices are new the vertex will be pushed to the end of the buffer.
	const auto new_vertex_index = static_cast<z3d::vertex_index>(m_mesh.positions().size());

	// Search through lookup to check if index combination already exists.
	const auto [ index_it, is_new ] = m_vertex_comp_indices_to_vertex_index.try_emplace(
		vertex_comp_indices,
		new_vertex_index
	);

	if (is_new)
	{
		const auto& [ position_index, tex_coord_index, normal_index ] = vertex_comp_indices;

		// If index is out of range, push default constructed value.
		// Not ideal, but better than out of range indices.
		// TODO let user at least know that something went wrong.

		if (position_index)
		{
			auto& position = m_mesh.positions().emplace_back();
			if (position_index <= m_position_buffer.size())
			{
				position = m_position_buffer[position_index - 1];
			}
		}

		if (normal_index)
		{
			auto& normal = m_mesh.normals().emplace_back();
			if (normal_index <= m_normal_buffer.size())
			{
				normal = m_normal_buffer[normal_index - 1];
			}
		}

		if (tex_coord_index)
		{
			auto& tex_coord = m_mesh.tex_coords().emplace_back();
			if (tex_coord_index <= m_tex_coord_buffer.size())
			{
				tex_coord = m_tex_coord_buffer[tex_coord_index - 1];
			}
		}
	}

	return index_it->second;
}

assets::obj_parser::prefetcher_context::prefetcher_context::prefetcher_context(
	path_id_lookups& id_lookups
) : m_id_lookups{ &id_lookups }
{
	m_buffer.resize(8192);
}

void assets::obj_parser::prefetcher_context::operator()(lookup_type::const_pointer entry) noexcept
{
	namespace fs = std::filesystem;

	const auto& [ filename, id ] = *entry;

	const auto base_dir = fs::canonical(fs::path(filename).parent_path());

	const auto buffer_view = std::string_view(m_buffer);

	auto filename_buffer = std::filesystem::path{};
	auto path_buffer = std::filesystem::path{};

	// Add linebreak to simplify line begin search.
	m_buffer.front() = '\n';
	auto leftover = std::size_t{ 1 };

	auto in = std::ifstream{ filename };

	static constexpr auto keyword = std::string_view{ "\nmtllib " };

	do
	{
		// Keep some old characters to continue matching interrupted sequence.
		std::copy(m_buffer.end() - leftover, m_buffer.end(), m_buffer.begin());

		in.read(m_buffer.data() + leftover, m_buffer.size() - leftover);

		const auto str = buffer_view.substr(0, leftover + in.gcount());

		leftover = keyword.size();

		auto pos = std::string_view::size_type{};
		while ((pos = str.find(keyword, pos)) != std::string_view::npos)
		{
			const auto filename_begin = pos + keyword.size();
			const auto filename_end = str.find('\n', filename_begin);

			if (filename_end != std::string_view::npos)
			{
				const auto length = filename_end - filename_begin;

				filename_buffer = str.substr(filename_begin, length);

				if (filename_buffer.is_relative())
				{
					path_buffer = base_dir;
					path_buffer /= filename_buffer;
				}

				m_id_lookups->material_libraries.try_emplace(path_buffer);

				pos = filename_end;
			}
			else // String match exceeds buffer.
			{
				if (pos == 0) [[unlikely]]
				{
					ztu::logger::error("Ignoring string match, as it exceeds buffer size of % characters.", m_buffer.size());
					leftover = 0;
				}
				else
				{
					leftover = str.size() - pos;
				}
				break;
			}
		}
	} while (not in.eof());
}