#include "viewer/asset_loader.hpp"

#include <geometry/normal_estimation.hpp>

#include "geometry/aabb.hpp"
#include "util/logger.hpp"
#include "glm/gtx/string_cast.hpp" // TODO remove

namespace viewer
{
std::error_code asset_loader::init(
	mesh_vertex_components::flags enabled_mesh_components,
	material_component::flags enabled_material_components,
	point_cloud_vertex_components::flags enabled_point_cloud_components,
	const dynamic_material_data& default_material
) {

	//m_ctx.setActive(true);

	m_enabled_mesh_components = enabled_mesh_components;
	m_enabled_material_components = enabled_material_components;
	m_enabled_point_cloud_components = enabled_point_cloud_components;

	m_dynamic_material_data_buffer.push_back(default_material);

	return create_gl_materials();
}

std::error_code asset_loader::load_shader(
	const GLenum type,
	const std::filesystem::path& filename,
	zgl::shader_handle& shader_handle
) {

	auto& [ buffer_source, buffer_type ] = m_dynamic_shader_data_buffer;

	std::stringstream source_stream;
	static constexpr auto glsl_version = 460;

	source_stream << "#version " << glsl_version << '\n';
	source_stream << "#define " << "hello" << '\n';
	source_stream << buffer_source;

	buffer_type = type;

	if (filename.empty())
	{
		buffer_source = "";
	}
	else
	{
		if (const auto e = m_shader_loader.load(filename, buffer_source)) {
			ztu::logger::warn(
				"Could not load shader_program_data source file %: [%] %",
				filename,
				e.category().name(),
				e.message()
			);
		}
	}

	if (const auto e = create_gl_shader()) {
		return e;
	}

	shader_handle = m_gl_shader_data.back().handle();

	return {};
}

std::error_code asset_loader::build_shader_program(
	const zgl::shader_handle& vertex_shader,
	const zgl::shader_handle& geometry_shader,
	const zgl::shader_handle& fragment_shader,
	zgl::shader_program_handle& shader_program_handle
) {

	auto program_data = zgl::shader_program_data{};

	if (const auto e = zgl::shader_program_data::build_from(
		vertex_shader,
		geometry_shader,
		fragment_shader,
		program_data
	)) {
		return e;
	}

	shader_program_handle = program_data.handle();

	m_gl_shader_program_data.emplace_back(std::move(program_data));

	return {};
}

std::error_code asset_loader::load_asset(
	const std::string& format,
	const std::filesystem::path& filename,
	std::vector<std::pair<dynamic_mesh_handle_type, dynamic_material_handle_type>>& dynamic_mesh_handles,
	std::vector<dynamic_point_cloud_handle_type>& dynamic_point_cloud_handles
) {
	std::error_code error;

	if ((error = load_mesh(
		format, filename, dynamic_mesh_handles
	))) {
		if (
			error.category() == std::generic_category() and
			static_cast<std::errc>(error.value()) == std::errc::invalid_argument
		) {
			error = load_point_cloud(format, filename, dynamic_point_cloud_handles);
		}
	}

	return error;
}

std::error_code asset_loader::load_mesh(
	const std::string& format,
	const std::filesystem::path& filename,
	std::vector<std::pair<dynamic_mesh_handle_type, dynamic_material_handle_type>>& dynamic_mesh_handles
) {
	const auto mesh_loader_id = m_mesh_loader.find_loader(format);

	if (not mesh_loader_id)
	{
		return std::make_error_code(std::errc::invalid_argument);
	}

	if (const auto e = m_mesh_loader.read(
		*mesh_loader_id,
		filename,
		m_dynamic_mesh_data_buffer,
		m_enabled_mesh_components,
		m_dynamic_material_data_buffer,
		m_enabled_material_components,
		next_materials_id,
		true
	)) {
		return e;
	}

	return process_materials_and_meshes(dynamic_mesh_handles);
}

std::error_code asset_loader::load_point_cloud(
	const std::string& format,
	const std::filesystem::path& filename,
	std::vector<dynamic_point_cloud_handle_type>& dynamic_point_cloud_handles
) {
	const auto point_cloud_loader_id = m_point_cloud_loader.find_loader(format);

	if (not point_cloud_loader_id)
	{
		return std::make_error_code(std::errc::invalid_argument);
	}

	if (const auto e = m_point_cloud_loader.read(
		*point_cloud_loader_id,
		filename,
		m_dynamic_point_cloud_buffer,
		true
	)) {
		return e;
	}

	return process_point_clouds(dynamic_point_cloud_handles);
}

std::error_code asset_loader::load_asset_directory(
	const std::string& format,
	const std::filesystem::path& path,
	std::vector<std::pair<dynamic_mesh_handle_type, dynamic_material_handle_type>>& dynamic_mesh_handles,
	std::vector<dynamic_point_cloud_handle_type>& dynamic_point_cloud_handles
) {
	std::error_code error;

	if ((error = load_mesh_directory(
		format, path, dynamic_mesh_handles
	))) {
		if (
			error.category() == std::generic_category() and
			static_cast<std::errc>(error.value()) == std::errc::invalid_argument
		) {
			error = load_point_cloud_directory(format, path, dynamic_point_cloud_handles);
		}
	}

	return error;
}
std::error_code asset_loader::load_mesh_directory(
	const std::string& format,
	const std::filesystem::path& path,
	std::vector<std::pair<dynamic_mesh_handle_type, dynamic_material_handle_type>>& dynamic_mesh_handles
) {
	const auto mesh_loader_id = m_mesh_loader.find_loader(format);

	if (not mesh_loader_id)
	{
		return std::make_error_code(std::errc::invalid_argument);
	}

	if (const auto e = m_mesh_loader.read_directory(
		*mesh_loader_id,
		path,
		m_dynamic_mesh_data_buffer,
		m_enabled_mesh_components,
		m_dynamic_material_data_buffer,
		m_enabled_material_components,
		next_materials_id,
		true
	)) {
		return e;
	}

	return process_materials_and_meshes(dynamic_mesh_handles);
}


std::error_code asset_loader::load_point_cloud_directory(
	const std::string& format,
	const std::filesystem::path& path,
	std::vector<dynamic_point_cloud_handle_type>& dynamic_point_cloud_handles
) {
	const auto point_cloud_loader_id = m_point_cloud_loader.find_loader(format);

	if (not point_cloud_loader_id)
	{
		return std::make_error_code(std::errc::invalid_argument);
	}

	if (const auto e = m_point_cloud_loader.read_directory(
		*point_cloud_loader_id,
		path,
		m_dynamic_point_cloud_buffer,
		true
	)) {
		return e;
	}

	return process_point_clouds(dynamic_point_cloud_handles);
}

std::error_code asset_loader::process_materials_and_meshes(
	std::vector<std::pair<dynamic_mesh_handle_type, dynamic_material_handle_type>>& dynamic_mesh_handles
) {

	const auto material_count_before = m_gl_material_data_references.size();

	if (const auto e = create_gl_materials())
	{
		m_dynamic_mesh_data_buffer.clear();
		return e;
	}

	const auto new_materials = std::span(
		m_gl_material_data_references.begin() + material_count_before,
		m_gl_material_data_references.end()
	);

	const auto mesh_count_before = m_gl_mesh_data.size();

	create_gl_meshes(new_materials);

	const auto new_meshes = std::span(
		m_gl_mesh_data.begin() + mesh_count_before,
		m_gl_mesh_data.end()
	);

	const auto dynamic_mesh_count_before = dynamic_mesh_handles.size();
	dynamic_mesh_handles.resize(dynamic_mesh_handles.size() + new_meshes.size());

	std::ranges::transform(
		new_meshes,
		dynamic_mesh_handles.begin() + dynamic_mesh_count_before,
		[&](const auto& entry)
		{
			const auto& [ gl_mesh_data, bounding_box ] = entry;
			const auto material_id = gl_mesh_data.material_id();

			auto material_index = std::size_t{ 0 };

			if (material_id != 0)
			{
				const auto material_reference_it = std::ranges::find_if(
					new_materials,
					[&material_id](const auto& entry) {
						return entry.first == material_id;
					}
				);

				if (material_reference_it == new_materials.end())
				{
					ztu::logger::error(
						"Something went horribly wrong while searching for material. Falling back to default material"
					);
				}
				else
				{
					material_index = material_reference_it.base() - m_gl_material_data_references.begin().base();
				}
			}

			const auto& gl_material = m_gl_material_data[material_index];

			//ztu::logger::debug("mesh components: %", static_cast<unsigned int>(gl_mesh_data.components()));
			//ztu::logger::debug("material components: %", static_cast<unsigned int>(gl_material.components()));

			return std::make_pair(
				dynamic_mesh_handle_type{
					.handle = gl_mesh_data.handle(),
					.bounding_box = bounding_box,
					.components = gl_mesh_data.components()
				},
				dynamic_material_handle_type{
					.handle = gl_material.handle(),
					.components = gl_material.components()
				}
			);
		}
	);

	return {};
}

std::error_code asset_loader::process_point_clouds(
	std::vector<dynamic_point_cloud_handle_type>& dynamic_point_cloud_handles
) {
	const auto point_cloud_count_before = m_gl_point_cloud_data.size();

	create_gl_point_clouds();

	const auto new_point_clouds = std::span(
		m_gl_point_cloud_data.begin() + point_cloud_count_before,
		m_gl_point_cloud_data.end()
	);

	const auto dynamic_point_cloud_count_before = dynamic_point_cloud_handles.size();
	dynamic_point_cloud_handles.resize(dynamic_point_cloud_handles.size() + new_point_clouds.size());

	std::ranges::transform(
		new_point_clouds,
		dynamic_point_cloud_handles.begin() + dynamic_point_cloud_count_before,
		[&](const auto& gl_point_cloud_data)
		{
			const auto& [ data, box ] = gl_point_cloud_data;
			return dynamic_point_cloud_handle_type{
				.handle = data.handle(),
				.bounding_box = box,
				.components = data.components()
			};
		}
	);

	return {};
}

std::error_code asset_loader::create_gl_materials()
{
	auto error = std::error_code{};

	for (const auto& material_data : m_dynamic_material_data_buffer)
	{
		auto gl_material_data = zgl::material_data{};

		if ((error = zgl::material_data::build_from(
			material_data.texture(),
			material_data.surface_properties(),
			material_data.transparency(),
			material_data.components(),
			gl_material_data
		))) {
			ztu::logger::error(
				"Error while creating material gpu handle: [%] %",
				error.category().name(),
				error.message()
			);
		}
		else
		{
			m_gl_material_data.emplace_back(std::move(gl_material_data));
			m_gl_material_data_references.emplace_back(next_materials_id, 0);
		}

		++next_materials_id;
	}

	m_dynamic_material_data_buffer.clear();

	return error;
}


void asset_loader::create_gl_meshes(std::span<const material_reference_entry_type> material_references)
{
	auto component_type_buffer = std::array<GLenum, static_cast<std::size_t>(mesh_vertex_components::count)>{};
	auto component_length_buffer = std::array<GLint, static_cast<std::size_t>(mesh_vertex_components::count)>{};

	auto component_stride = GLsizei{};
	auto component_count = ztu::u32{};

	for (auto& mesh_data :  m_dynamic_mesh_data_buffer)
	{
		if (mesh_data.triangles().empty())
		{
			ztu::logger::warn("Skipping mesh with empty index buffer.");
			continue;
		}

		const auto material_id = mesh_data.material_id();

		auto material_index = std::size_t{ 0 };

		if (material_id != 0) // Default material is always there
		{
			const auto material_reference_it = std::ranges::find_if(
				material_references,
				[&material_id](const material_reference_entry_type& entry) {
					return entry.first == material_id;
				}
			);

			if (material_reference_it == material_references.end())
			{
				ztu::logger::error("Skipping mesh because referenced material cannot be found");
				continue;
			}

			material_index = material_reference_it - material_references.begin();
		}

		// Add normals if missing
		if ((mesh_data.components() & mesh_vertex_components::flags::normal) == mesh_vertex_components::flags::none)
		{
			ztu::logger::warn("Model is missing normal vectors, so they are estimated!");
			estimate_normals(
				mesh_data.positions(),
				mesh_data.triangles(),
				mesh_data.normals()
			);
			mesh_data.components() |= mesh_vertex_components::flags::normal;
		}

		auto mesh_box = aabb{};
		mesh_box.add_points<mesh_vertex_components::normal::value_type>(mesh_data.positions());

		mesh_data.build_vertex_buffer(
			m_vertex_buffer,
			component_count,
			component_type_buffer,
			component_length_buffer,
			component_stride
		);

		auto gl_mesh_data = zgl::mesh_data{};

		const auto& first_triangle = mesh_data.triangles().front();

		// TODO make span of size component_count

		if (const auto e = zgl::mesh_data::build_from(
			m_vertex_buffer,
			std::span(component_type_buffer).subspan(0, component_count),
			std::span(component_length_buffer).subspan(0, component_count),
			component_stride,
			std::span(
				first_triangle.data(),
				mesh_data.triangles().size() * first_triangle.size()
			),
			mesh_data.material_id(),
			mesh_data.components(),
			gl_mesh_data
		)) {
			ztu::logger::error(
				"Error while creating opengl mesh data: [%] %\nMesh will be skipped.",
				e.category().name(),
				e.message()
			);
		}

		++m_gl_material_data_references[material_index].second;

		m_gl_mesh_data.emplace_back(std::move(gl_mesh_data), mesh_box);
	}

	m_dynamic_mesh_data_buffer.clear();
}

void asset_loader::create_gl_point_clouds()
{
	auto component_type_buffer = std::array<GLenum, static_cast<std::size_t>(point_cloud_vertex_components::count)>{};
	auto component_length_buffer = std::array<GLint, static_cast<std::size_t>(point_cloud_vertex_components::count)>{};
	auto component_stride = GLsizei{};
	auto component_count = ztu::u32{};

	for (const auto& point_cloud_data : m_dynamic_point_cloud_buffer)
	{
		if (point_cloud_data.positions().empty())
		{
			ztu::logger::warn("Skipping point cloud without points.");
			continue;
		}

		auto point_cloud_box = aabb{};
		point_cloud_box.add_points<mesh_vertex_components::normal::value_type>(point_cloud_data.positions());

		point_cloud_data.build_vertex_buffer(
			m_vertex_buffer,
			component_count,
			component_type_buffer,
			component_length_buffer,
			component_stride
		);

		auto gl_point_cloud_data = zgl::point_cloud_data{};

		if (const auto e = zgl::point_cloud_data::build_from(
			m_vertex_buffer,
			component_type_buffer,
			component_length_buffer,
			component_stride,
			gl_point_cloud_data
		)) {
			ztu::logger::error(
				"Error while creating opengl point cloud data: [%] %\nPoint cloud will be skipped.",
				e.category().name(),
				e.message()
			);
		}

		m_gl_point_cloud_data.emplace_back(std::move(gl_point_cloud_data), point_cloud_box);
	}

	m_dynamic_point_cloud_buffer.clear();
}

std::error_code asset_loader::create_gl_shader()
{
	auto shader_data = zgl::shader_data{};

	const auto& [source, type] = m_dynamic_shader_data_buffer;

	if (const auto e = zgl::shader_data::build_from(type, source, shader_data)) {
		return e;
	}

	m_gl_shader_data.emplace_back(std::move(shader_data));


	return {};
}


bool asset_loader::unload(const zgl::shader_program_handle& shader_handle)
{
	const auto it = std::find_if(
		m_gl_shader_program_data.begin(), m_gl_shader_program_data.end(),
		[&shader_handle](const auto& gl_shader_data) {
			return gl_shader_data.handle().program_id == shader_handle.id;
		}
	);

	if (it == m_gl_shader_program_data.end())
	{
		return false;
	}

	m_gl_shader_program_data.erase(it);

	return true;
}

bool asset_loader::unload(const zgl::mesh_handle& mesh_handle)
{
	const auto it = std::ranges::find_if(
		m_gl_mesh_data,
		[&mesh_handle](const auto& gl_mesh_data) {
			return gl_mesh_data.first.handle().vao_id == mesh_handle.vao_id;
		}
	);

	if (it == m_gl_mesh_data.end())
	{
		return false;
	}

	const auto material_id = it->first.material_id();

	const auto reference_it = std::ranges::find_if(
		m_gl_material_data_references,
		[&material_id](const auto& entry) {
			return entry.first == material_id;
		}
	);

	if (reference_it != m_gl_material_data_references.end())
	{
		// Do not delete default material at index 0
		if (--reference_it->second == 0 and reference_it != m_gl_material_data_references.begin())
		{
			const auto index = reference_it - m_gl_material_data_references.begin();
			m_gl_material_data.erase(m_gl_material_data.begin() + index);
			m_gl_material_data_references.erase(reference_it);
		}
	}

	m_gl_mesh_data.erase(it);

	return true;
}

void asset_loader::unload_shader_data()
{
	m_gl_shader_data.clear();
}


bool asset_loader::unload(const zgl::point_cloud_handle& point_cloud_handle)
{
	const auto it = std::ranges::find_if(
		m_gl_point_cloud_data,
		[&point_cloud_handle](const auto& entry) {
			return entry.first.handle().vao_id == point_cloud_handle.vao_id;
		}
	);

	if (it == m_gl_point_cloud_data.end())
	{
		return false;
	}

	m_gl_point_cloud_data.erase(it);

	return true;
}
}