#include "assets/file_parsers/obj_loader.hpp"

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

#include "assets/components/mesh_vertex_components.hpp"
#include "assets/data_loaders/"

#include "util/logger.hpp"
#include "util/for_each.hpp"
#include "util/uix.hpp"
#include <set>

#include "util/line_parser.hpp"

namespace obj_loader_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& connector_error_category()
{
	static obj_loader_error::category category;
	return category;
}

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

template<typename T, std::size_t Count>
std::errc parse_numeric_vector(std::string_view param, std::array<T, Count>& values)
{
	auto it = param.begin();
	const auto end = param.end();

	for (auto& value : values)
	{
		if (it >= end)
		{
			return std::errc::invalid_argument;
		}

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

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

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

	return {};
};


void obj_loader::find_materials(
	std::span<char> buffer,
	std::filesystem::path& path_buffer,
	const std::filesystem::path& base_directory,
	std::ifstream& in,
	ztu::string_list& material_filenames
) {
	static constexpr auto keyword = std::string_view{ "\nmtllib " };

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

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

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

		in.read(buffer.data() + leftover, 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;
				const auto filename = str.substr(filename_begin, length);

				path_buffer.assign(filename);

				if (path_buffer.is_relative())
				{
					path_buffer = base_directory;
					path_buffer /= filename;
				}

				material_filenames.push_back(path_buffer.c_str());

				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.", buffer.size());
					leftover = 0;
				}
				else
				{
					leftover = str.size() - pos;
				}
				break;
			}
		}
	} while (not in.eof());
}


std::error_code obj_loader::prefetch(
	const file_dir_list& paths,
	prefetch_queue& queue
) {
	namespace fs = std::filesystem;
	using obj_loader_error::codes;
	using obj_loader_error::make_error_code;

	auto buffer = std::vector<char>(8 * 1024, '\0');

	auto in = std::ifstream{};

	auto path_buffer = fs::path{};
	auto filename_buffer = fs::path{};

	const auto process_file = [&]()
	{
		in.open(filename_buffer);
		if (not in.is_open()) {
			ztu::logger::error("Could not open .obj file '%'", filename_buffer);
			return;
		}

		filename_buffer.remove_filename();

		find_materials(buffer, path_buffer, filename_buffer, in, queue.mtl_queue.files);

		in.close();
	};


	for (const auto file : paths.files)
	{
		filename_buffer.assign(file);
		process_file();
	}

	for (const auto directory : paths.directories)
	{
		for (const auto& file : fs::directory_iterator{ directory }) {

			filename_buffer.assign(file.path());

			// Avoid heap allocation of .extension()
			if (not std::string_view(filename_buffer.c_str()).ends_with(".obj"))
			{
				continue;
			}

			process_file();
		}
	}

	return {};
}

std::error_code obj_loader::load(
	dynamic_mesh_buffer& buffer,
	const file_dir_list& paths,
	prefetch_lookup& id_lookup,
	dynamic_shader_source_store& store,
	bool pedantic
) {
	namespace fs = std::filesystem;

	auto position_buffer = buffer.positions();
	auto normal_buffer = buffer.normals();
	auto tex_coord_buffer = buffer.tex_coords();

	auto read_buffer = dynamic_mesh_buffer{};
	auto path_buffer = fs::path{};
	auto vertex_ids = std::set<indexed_vertex_type>{};
	auto in = std::ifstream{};

	auto filename_buffer = fs::path{};

	const auto process_file = [&]()
	{
		in.open(filename_buffer);
		if (not in.is_open()) {
			ztu::logger::error("Could not open .obj file '%'", filename_buffer);
			return;
		}

		filename_buffer.remove_filename();

		// parse file
		const auto error = parse_file(
			read_buffer,
			buffer,
			path_buffer,
			filename_buffer,
			vertex_ids,
			in,
			id_lookup,
			store,
			pedantic
		);

		if (error)
		{
			ztu::logger::error(
				"Error occurred while parsing .obj file: [%] %",
				error.category().name(),
				error.message()
			);
		}

		in.close();
	};


	for (const auto file : paths.files)
	{
		filename_buffer.assign(file);
		process_file();
	}

	for (const auto directory : paths.directories)
	{
		for (const auto& file : fs::directory_iterator{ directory }) {

			filename_buffer.assign(file.path());

			// Avoid heap allocation of .extension()
			if (not std::string_view(filename_buffer.c_str()).ends_with(".obj"))
			{
				continue;
			}

			process_file();
		}
	}

	return {};
}


std::error_code obj_loader::parse_file(
	dynamic_mesh_buffer& read_buffer,
	dynamic_mesh_buffer& mesh_buffer,
	std::filesystem::path& path_buffer,
	const std::filesystem::path& base_directory,
	std::set<indexed_vertex_type>& vertex_ids,
	std::ifstream& in,
	prefetch_lookup& id_lookup,
	dynamic_shader_source_store& store,
	bool pedantic
) {
	using obj_loader_error::codes;
	using obj_loader_error::make_error_code;

	read_buffer.clear();
	mesh_buffer.clear();
	vertex_ids.clear();

	// Buffers for storing the vertex component definitions.
	auto& position_buffer = read_buffer.positions();
	auto& normal_buffer = read_buffer.normals();
	auto& tex_coord_buffer = read_buffer.tex_coords();

	auto& positions = mesh_buffer.positions();
	auto& normals = mesh_buffer.normals();
	auto& tex_coords = mesh_buffer.tex_coords();
	auto& triangles = mesh_buffer.triangles();

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

			// Copy buffer into store and keep capacity.
			store.meshes.add(mesh_buffer);
		}

		if (clear_read_buffer)
		{
			read_buffer.clear();
		}

		mesh_buffer.clear();
		vertex_ids.clear();
	};

	const auto find_or_push_vertex = [&](const vertex_type& vertex) -> index_type
	{
		auto indexed_vid = indexed_vertex_type{
			.vertex = vertex,
			.buffer_index = static_cast<index_type>(positions.size())
		};

		// Search through sorted lookup to check if index combination is unique
		const auto [ id_it, unique ] = vertex_ids.insert(indexed_vid);

		if (unique)
		{
			const auto& [ position_index, tex_coord_index, normal_index ] = vertex;

			// If index is out of range, push default constructed value.
			// Not ideal, but better than out of range indices.

			auto& position = positions.emplace_back();
			if (position_index < position_buffer.size())
			{
				position = position_buffer[position_index];
			}

			auto& normal = normals.emplace_back();
			if (normal_index < normal_buffer.size())
			{
				normal = normal_buffer[normal_index];
			}

			auto& tex_coord = tex_coords.emplace_back();
			if (tex_coord_index < tex_coord_buffer.size())
			{
				tex_coord = tex_coord_buffer[tex_coord_index];
			}
		}

		return id_it->buffer_index;
	};

	auto curr_material_library_it = dynamic_material_library_store::iterator_type{};
	auto has_material_library = false;

	const auto ec = ztu::parse_lines<codes>(
		in,
		pedantic,
		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;
			}

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

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

			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();

			auto vertex = vertex_type{};

			index_type first_index{}, prev_index{};

			auto vertex_count = std::size_t{};

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

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

					--component_index; // Convert to zero based index.
					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)
				{
					 auto& triangle = triangles.emplace_back();
					triangle[0] = first_index;
					triangle[1] = prev_index;
					triangle[2] = curr_index;
				}
				else if (vertex_count == 1)
				{
					first_index = curr_index;
				}

				prev_index = curr_index;
		   }

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

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

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

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

			mesh_buffer.material_id() = 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_directory;
				path_buffer /= param; // TODO Doesn't thi allocate an extra path?!?
			}

			const auto material_library_id_it = id_lookup.material_libraries.find(path_buffer);

			if (material_library_id_it != id_lookup.material_libraries.end()) [[likely]]
			{
				const auto material_library_id = material_library_id_it->second;
				std::tie(curr_material_library_it, has_material_library) = store.material_libraries.find(material_library_id);
			}
			else [[unlikely]]
			{
				ztu::logger::warn(
					"Could not find a matching material library with path '%'. Proceeding with default material.",
					param
				);
				has_material_library = false;
			}
		})
	);

	if (ec != codes::ok)
	{
		return make_error_code(ec);
	}

	push_mesh();

	return {};
}