#include "assets/data_loaders/stl_loader.hpp"

#include "util/binary_ifstream.hpp"
#include "util/unroll_bool_template.hpp"
#include "util/logger.hpp"

template<bool Normals>
std::error_code read_body(
	binary_ifstream& in,
	const std::uint32_t expected_triangle_count,
	std::vector<components::mesh_vertex::position>& positions,
	std::vector<components::mesh_vertex::normal>& normals,
	std::vector<std::array<ztu::u32, 3>>& triangles
) {

	const auto read_vector = [&in](auto& vector) -> std::error_code
	{
		for (auto& component : vector)
		{
			float component32;
			if (const auto e = in.read_ieee754<std::endian::little>(component32))
			{
				return e;
			}
			component = component32;
		}
		return {};
	};

	for (std::uint32_t i{}; i != expected_triangle_count; ++i) {

		auto normal = components::mesh_vertex::normal{};
		if constexpr (Normals)
		{
			if (const auto e = read_vector(normal))
			{
				return e;
			}
		}

		auto triangle = std::array<ztu::u32, 3>{};

		for (auto& index : triangle) {

			auto position = components::mesh_vertex::position{};
			if (const auto e = read_vector(position))
			{
				return e;
			}

			// TODO implement unique insert correctly
			/*// Insert vertices sorted, to efficiently remove duplicates.
			const auto it = std::ranges::upper_bound(positions, position);

			// Set index before `it` is invalidated by insert.
			index = it - positions.begin();

			if (it != positions.begin() and *std::prev(it) == position)
			{
				--index;
			}
			else
			{
				positions.insert(it, position);
				if constexpr (Normals)
				{
					normals.insert(normals.begin() + index, normal);
				}
			}*/

			index = positions.size();
			positions.push_back(position);

			if constexpr (Normals)
			{
				normals.push_back(normal);
			}
		}

		triangles.push_back(triangle);

		// Skip attribute bytes
		if (const auto e = in.skip<std::uint16_t>())
		{
			return e;
		}
	}

	return {};
}

std::error_code stl_loader::read_directory(
	const std::filesystem::path& path,
	std::vector<dynamic_mesh_data>& meshes,
	components::mesh_vertex::flags enabled_components::mesh_vertexs,
	std::vector<dynamic_material_data>& materials,
	material_component::flags enabled_material_components,
	const ztu::u32 base_material_id,
	bool pedantic
) {
	namespace fs = std::filesystem;

	if (not fs::exists(path)) {
		return make_error_code(std::errc::no_such_file_or_directory);
	}

	for (const auto& file : fs::directory_iterator{ path / "frames" })
	{
		const auto& file_path = file.path();

		if (file_path.extension() != ".stl")
		{
			continue;
		}

		if (const auto e = read(
			file_path,
			meshes,
			enabled_components::mesh_vertexs,
			materials,
			enabled_material_components,
			base_material_id,
			pedantic
		)) {
			ztu::logger::error(
				"Error while loading stl file '%': [%] %",
				file_path,
				e.category().name(),
				e.message()
			);
		}
	}

	return {};
}


std::error_code stl_loader::read(
	const std::filesystem::path& filename,
	std::vector<dynamic_mesh_data>& meshes,
	components::mesh_vertex::flags enabled_components::mesh_vertexs,
	std::vector<dynamic_material_data>&,
	material_component::flags,
	ztu::u32,
	const bool pedantic
) {
	auto error = std::error_code{};

	auto in = binary_ifstream{};

	if ((error = in.open(filename, true)))
	{
		return error;
	}

	auto header_bytes_left = static_cast<binary_ifstream::size_type>(80);

	if (pedantic)
	{
		// Check if ASCII file was provided, these start with a specific character sequence.

		static constexpr auto ascii_magic_string = std::string_view("solid");

		auto magic_bytes = std::array<binary_ifstream::char_type, ascii_magic_string.size()>{};

		if ((error = in.read(magic_bytes)))
		{
			return error;
		}

		const auto magic_string = std::string_view(
			reinterpret_cast<const char*>(magic_bytes.data()),
			magic_bytes.size()
		);

		if (magic_string == ascii_magic_string)
		{
			return std::make_error_code(std::errc::illegal_byte_sequence);
		}

		header_bytes_left -= ascii_magic_string.size();
	}

	// Ignore (rest of) header.
	if ((error = in.skip(header_bytes_left)))
	{
		return error;
	}

	// Read number of bytes
	auto expected_triangle_count = std::uint32_t{};

	if ((error = in.read<std::endian::little>(expected_triangle_count)))
	{
		return error;
	}

	// Use separate mesh for parsing, so original mesh is only overwritten
	// if no errors occurred. This also guarantees unused reserved memory
	// is freed immediately in case of an error.
	auto mesh = dynamic_mesh_data{};

	auto& positions = mesh.positions();
	auto& normals = mesh.normals();
	auto& triangles = mesh.triangles();
	auto& material_id = mesh.material_id();

	material_id = 0; // Set to default material

	positions.reserve(expected_triangle_count * 3);
	normals.reserve(expected_triangle_count);
	triangles.reserve(expected_triangle_count);

	const auto normals_enabled = (
		(enabled_components::mesh_vertexs & components::mesh_vertex::flags::normal) != components::mesh_vertex::flags::none
	);

	error = unroll_bool_function_template([&]<bool Normals>() {
		return read_body<Normals>(
			in,expected_triangle_count,
			positions,
			normals,
			triangles
		);
	}, normals_enabled);

	// Free any unused reserved memory
	positions.shrink_to_fit();
	normals.shrink_to_fit();
	triangles.shrink_to_fit();

	if (error)
	{
		return error;
	}

	ztu::logger::debug("Normal count: %", normals.size());

	if (not positions.empty())
	{
		mesh.components() |= components::mesh_vertex::flags::position;
	}

	if (not normals.empty())
	{
		ztu::logger::debug("Enabling normals!!!");
		mesh.components() |= components::mesh_vertex::flags::normal;
	}

	meshes.emplace_back(std::move(mesh));

	return {};
}