#include <psemek/gfx/gltf_mesh.hpp>
#include <psemek/gfx/gltf_accessor_iterator.hpp>
#include <psemek/gfx/array.hpp>
#include <psemek/gfx/buffer.hpp>
#include <psemek/util/hstring.hpp>
#include <psemek/util/hash_table.hpp>
#include <psemek/io/memory_stream.hpp>

namespace psemek::gfx
{

	namespace
	{

		struct drawable_impl
			: drawable
		{
			gfx::array vao;
			std::size_t index_count;
			std::size_t index_offset;
			GLenum index_type;

			std::vector<math::point<float, 3>> vertices;
			std::vector<math::triangle<std::uint32_t>> triangles;

			void draw() const override
			{
				vao.bind();
				gl::DrawElements(gl::TRIANGLES, index_count, index_type, reinterpret_cast<void const *>(index_offset));
			}
		};

		struct gltf_mesh_impl
			: gltf_mesh
		{
			gltf_mesh_impl(gltf_asset const & asset, std::function<util::blob(std::string const &)> uri_loader);

			gltf_asset::material const & material(std::size_t index) const override
			{
				return materials_[index];
			}

			gfx::texture_2d const & texture(std::size_t index) const override
			{
				return textures_[index];
			}

			util::span<primitive const> mesh(std::string_view name) const override
			{
				if (auto it = meshes_.find(name); it != meshes_.end())
					return it->second.primitives;
				return {};
			}

			gltf_asset::extra const * mesh_property(std::string_view mesh_name, std::string_view property_name) const override
			{
				if (auto it = meshes_.find(mesh_name); it != meshes_.end())
					if (auto jt = it->second.extras.find(property_name); jt != it->second.extras.end())
						return &(jt->second);
				return nullptr;
			}

		private:
			struct mesh_data
			{
				std::vector<primitive> primitives;
				gltf_asset::extras_map extras;
				std::vector<math::point<float, 3>> vertices;
				std::vector<math::triangle<std::uint32_t>> triangles;
			};

			std::vector<gltf_asset::material> materials_;
			std::vector<gfx::texture_2d> textures_;
			std::vector<gfx::buffer> buffers_;
			std::vector<std::unique_ptr<drawable_impl>> drawables_;
			util::hash_map<util::hstring, mesh_data> meshes_;
		};

		gltf_mesh_impl::gltf_mesh_impl(gltf_asset const & asset, std::function<util::blob(std::string const &)> uri_loader)
		{
			materials_ = asset.materials;

			for (auto const & buffer : asset.buffers)
			{
				if (buffer.data)
				{
					buffers_.emplace_back().load(buffer.data->data(), buffer.data->size(), gl::STATIC_DRAW);
				}
				else
				{
					auto data = uri_loader(buffer.uri);
					buffers_.emplace_back().load(data.data(), data.size(), gl::STATIC_DRAW);
				}
			}

			for (auto const & texture : asset.textures)
			{
				auto & target = textures_.emplace_back();

				if (texture.uri)
				{
					auto data = uri_loader(*texture.uri);
					target.load_srgb(gfx::read_image<gfx::color_rgba>(io::memory_istream(data.string_view())));
				}
				else
				{
					target.load_srgb(gfx::pixmap_rgba({2, 2}, gfx::color_rgba{255, 0, 255, 255}));
				}
				target.linear_mipmap_filter();
				target.anisotropy();
				target.generate_mipmap();
			}

			for (auto const & node : asset.nodes)
			{
				if (!node.mesh)
					continue;

				auto & target_mesh = meshes_[node.name];
				target_mesh.extras = node.extras;

				auto const & mesh = asset.meshes[*node.mesh];

				for (auto const & primitive : mesh.primitives)
				{
					drawables_.push_back(std::make_unique<drawable_impl>());
					auto & drawable = *drawables_.back();

					drawable.vao.bind();

					{
						auto const & indices_accessor = asset.accessors[primitive.indices];
						auto const & indices_view = asset.buffer_views[indices_accessor.buffer_view];

						gl::BindBuffer(gl::ELEMENT_ARRAY_BUFFER, buffers_[indices_view.buffer].id());
						drawable.index_offset = indices_view.offset;
						drawable.index_type = indices_accessor.component_type;
						drawable.index_count = indices_accessor.count;
					}

					std::pair<GLuint, std::optional<std::size_t>> attributes[] =
					{
						{0, primitive.position},
						{1, primitive.normal},
						{2, primitive.texcoord},
						{3, primitive.color},
					};

					for (auto const & attribute : attributes)
					{
						if (!attribute.second) continue;

						auto const & accessor = asset.accessors[*attribute.second];
						auto const & view = asset.buffer_views[accessor.buffer_view];

						gl::BindBuffer(gl::ARRAY_BUFFER, buffers_[view.buffer].id());
						gl::EnableVertexAttribArray(attribute.first);
						gl::VertexAttribPointer(attribute.first, attribute_size(accessor.type), accessor.component_type, accessor.normalized, 0, reinterpret_cast<void const *>(view.offset));
					}

					{
						auto indices = accessor_range<std::uint32_t>(asset, primitive.indices);
						for (auto it = indices.it_begin; it != indices.it_end;)
						{
							auto i0 = *it++;
							auto i1 = *it++;
							auto i2 = *it++;
							drawable.triangles.push_back({i0, i1, i2});
						}
					}

					if (primitive.position)
						for (auto const & p : accessor_range<math::point<float, 3>>(asset, *primitive.position))
							drawable.vertices.push_back(p);

					target_mesh.primitives.push_back({&drawable, primitive.material, drawable.vertices, drawable.triangles});
				}
			}
		}

	}

	std::unique_ptr<gltf_mesh> make_gltf_mesh(gltf_asset const & asset, std::function<util::blob(std::string const &)> uri_loader)
	{
		return std::make_unique<gltf_mesh_impl>(asset, std::move(uri_loader));
	}

}