#include "assets/file_parsers/kitti_parser.hpp"

#include "glm/glm.hpp"

#include <charconv>
#include <fstream>
#include <glm/ext/matrix_transform.hpp>

#include "assets/components/point_cloud_vertex_components.hpp"
#include "assets/data/pose_data.hpp"
#include "util/binary_ifstream.hpp"
#include "util/logger.hpp"
#include <execution>


assets::kitti_parser::parser_context::parser_context(
	const pose_list_id_lookup& pose_list_lookup,
	const pose_list_store& pose_list_store,
	store_type& m_store
) :
	m_pose_list_lookup{ &pose_list_lookup },
	m_pose_list_store{ &pose_list_store },
	m_store{ &m_store }
{
	constexpr auto expected_vertex_count = 8192;
	m_buffer.positions().reserve(expected_vertex_count);
	m_buffer.normals().reserve(expected_vertex_count);
	m_buffer.colors().reserve(expected_vertex_count);
}

void assets::kitti_parser::parser_context::reset()
{
	m_buffer.clear();
}

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

	pose_data pose;

	if (const auto pose_path = get_pose_path(filename); not pose_path)
	{
		if (pose_path != m_last_pose_path)
		{
			if (const auto pose_list_id_it = m_pose_list_lookup->find(*pose_path); pose_list_id_it != m_pose_list_lookup->end())
			{
				m_last_pose_path = *pose_path;
				const auto pose_list_id = pose_list_id_it->second;

				if (const auto [ pose_list_it, found ] = m_pose_list_store->find(pose_list_id); found)
				{
					m_last_pose_path = *pose_path;
					m_last_pose_list = pose_list_it->second;
				}
				else
				{
					ztu::logger::error("No matching pose found in store for %.", *pose_path);
					return;
				}
			}
			else
			{
				ztu::logger::error("No matching pose registered in lookup %.", *pose_path);
				return;
			}
		}
	}
	else
	{
		ztu::logger::error("Malformed kitti file path %.", filename);
		return;
	}

	reset();

	if (const auto e = load_point_file(filename, m_buffer))
	{
		ztu::logger::error("Could not load kitti file %: %", filename, e.message());
		return;
	}

	transform_point_cloud(m_buffer.positions(), pose);

	m_store->insert(id, m_buffer);
}

ztu::result<std::filesystem::path> assets::kitti_parser::parent_directory(
	const std::filesystem::path& path
) {
	try
	{
		return path.parent_path();
	}
	catch (const std::exception&)
	{
		return std::unexpected(std::make_error_code(std::errc::no_such_file_or_directory));
	}
};

ztu::result<std::filesystem::path> assets::kitti_parser::get_pose_path(
	const std::filesystem::path& path
) {
	return parent_directory(path)
		.and_then(parent_directory)
		.and_then(
			[](const auto& base_dir)
			{
				return base_dir / "pose.txt";
			}
		);
};

std::error_code assets::kitti_parser::prefetch(
	path_id_lookups& lookups
) {
	m_path_buffer.clear();
	lookups.point_clouds.by_extension(".bin", m_path_buffer);

	for (const auto entry : m_path_buffer)
	{

		if (const auto pose_path = pose_path(entry->second))
		{
			lookups.poses.try_emplace(std::filesystem::absolute(pose_path));
		}
		else
		{
			ztu::logger::error("Malformed kitti file path %.", entry->second);
		}
	}

	return {};
}

std::error_code assets::kitti_parser::load(
	 path_id_lookups& lookups,
	 data_stores& stores,
	 bool pedantic
) {
	namespace fs = std::filesystem;

	m_path_buffer.clear();
	lookups.point_clouds.by_extension(".bin", m_path_buffer);

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

	return {};
}

void assets::kitti_parser::transform_point_cloud(
	std::span<point_cloud_vertex_components::position> points,
	const glm::mat4& pose
) {
	for (auto& point : points)
	{
		point  = pose * glm::vec4{ point, 1.0f };
	}
}

std::error_code assets::kitti_parser::load_point_file(
	const std::filesystem::path& filename,
	point_cloud_data& point_cloud
) {

	auto in = binary_ifstream{};

	if (const auto e = in.open(filename, true))
	{
		return e;
	}

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

	point_cloud_vertex_components::position position;

	auto& positions = point_cloud.positions();

	auto parsing_error = std::error_code{};

	while (not ((parsing_error = read_vector(position))))
	{
		positions.push_back(position);

		if ((parsing_error = in.skip<float>())) // TODO what am I skipping here?!?
		{
			break;
		}
	}

	if (static_cast<std::errc>(parsing_error.value()) != std::errc::result_out_of_range)
	{
		return parsing_error;
	}

	return {};
}

ztu::result<std::size_t> assets::kitti_parser::frame_id_from_filename(
	const std::string_view filename
) {
	std::size_t id;
	const auto result = std::from_chars(filename.cbegin(), filename.cend(), id);

	if (result.ec != std::errc{})
	{
		return std::unexpected(std::make_error_code(result.ec));
	}

	return id;
}