#include <psemek/gfx/mesh.hpp>

#include <psemek/util/to_string.hpp>
#include <psemek/util/binary_stream.hpp>
#include <psemek/util/exception.hpp>

namespace psemek::gfx
{

	mesh::mesh(mesh && other)
		: array_{std::move(other.array_)}
		, vertex_buffer_{std::move(other.vertex_buffer_)}
		, index_buffer_{std::move(other.index_buffer_)}
		, instance_buffer_{std::move(other.instance_buffer_)}
		, info_{other.info_}
	{
		other.info_ = mesh_info{};
	}

	mesh & mesh::operator = (mesh && other)
	{
		if (this == &other) return *this;

		array_ = std::move(other.array_);
		vertex_buffer_ = std::move(other.vertex_buffer_);
		index_buffer_ = std::move(other.index_buffer_);
		instance_buffer_ = std::move(other.instance_buffer_);
		info_ = other.info_;
		other.info_ = mesh_info{};

		return *this;
	}

	void mesh::bind() const
	{
		array_.bind();
	}

	void mesh::setup(attribs_description const & attribs)
	{
		assert(!attribs.attribs.empty());

		bind();

		for (std::size_t i = 0; i < info_.max_attribute_index_; ++i)
			gl::DisableVertexAttribArray(i);

		std::size_t max_index = 0;
		bool instanced = false;

		for (auto const & a : attribs.attribs)
		{
			max_index = std::max<std::size_t>(max_index, a.index);

			if (a.divisor != 0)
			{
				assert(a.divisor == 1);
				instanced = true;
			}
		}

		info_.max_attribute_index_ = max_index;
		info_.stride_ = attribs.vertex_size;
		info_.instance_stride_ = attribs.instance_size;
		info_.instanced_ = instanced;

		if (!vertex_buffer_) vertex_buffer_ = buffer{};
		vertex_buffer_.bind();

		for (auto const & a : attribs.attribs)
		{
			if (a.divisor != 0) continue;

			gl::EnableVertexAttribArray(a.index);
			if (a.integer)
				gl::VertexAttribIPointer(a.index, a.size, a.type, attribs.vertex_size, a.pointer);
			else
				gl::VertexAttribPointer(a.index, a.size, a.type, a.normalized, attribs.vertex_size, a.pointer);
		}

		if (instanced)
		{
			if (!instance_buffer_) instance_buffer_ = buffer{};
			instance_buffer_.bind();

			for (auto const & a : attribs.attribs)
			{
				if (a.divisor == 0) continue;

				gl::EnableVertexAttribArray(a.index);
				gl::VertexAttribDivisor(a.index, a.divisor);
				if (a.integer)
					gl::VertexAttribIPointer(a.index, a.size, a.type, attribs.instance_size, a.pointer);
				else
					gl::VertexAttribPointer(a.index, a.size, a.type, a.normalized, attribs.instance_size, a.pointer);
			}
		}
	}

	mesh mesh::create(imported_mesh const & m)
	{
		mesh result;
		result.setup(m.attribs);
		result.load_raw(m);
		return result;
	}

	void mesh::load_raw(void const * vertices, std::size_t vertex_size, std::size_t count, GLenum primitive_type, GLenum usage)
	{
		assert(info_.stride_ == vertex_size);

		if (auto n = detail::get_primitive_type_vertex_count(primitive_type); n)
			assert((count % (*n)) == 0);

		assert(vertex_buffer_);

		vertex_buffer_.load(vertices, vertex_size * count, usage);
		info_.vertex_count_ = count;
		info_.index_count_ = 0;
		info_.indexed_ = false;
		info_.primitive_type_ = primitive_type;
	}

	static std::size_t index_size(GLenum type)
	{
		switch (type)
		{
		case gl::UNSIGNED_BYTE: return 1;
		case gl::UNSIGNED_SHORT: return 2;
		case gl::UNSIGNED_INT: return 4;
		default: throw util::exception("Unknown undex type");
		}
	}

	void mesh::load_raw(void const * vertices, std::size_t vertex_size, std::size_t vertex_count,
		void const * indices, GLenum index_type, std::size_t index_count,
		GLenum primitive_type, GLenum usage)
	{
		assert(info_.stride_ == vertex_size);

		if (auto n = detail::get_primitive_type_vertex_count(primitive_type); n)
			assert((index_count % (*n)) == 0);

		assert(vertex_buffer_);

		if (!index_buffer_)
			index_buffer_ = {};
		array_.bind();
		vertex_buffer_.load(vertices, vertex_size * vertex_count, usage);
		index_buffer_.load(indices, index_size(index_type) * index_count, usage);
		info_.vertex_count_ = vertex_count;
		info_.index_count_ = index_count;
		info_.indexed_ = true;
		info_.primitive_type_ = primitive_type;
		info_.index_type_ = index_type;
	}

	void mesh::load_raw(imported_mesh const & m)
	{
		setup(m.attribs);
		load_raw(m.vertices.data(), m.attribs.vertex_size, m.vertices.size() / m.attribs.vertex_size, m.indices.data(), gl::UNSIGNED_INT, m.indices.size(), gl::TRIANGLES, gl::STATIC_DRAW);
	}

	void mesh::draw() const
	{
		draw(0, is_indexed() ? index_count() : vertex_count(), instance_count());
	}

	void mesh::draw(std::size_t first, std::size_t count, std::size_t instance_count) const
	{
		if (is_indexed())
		{
			assert(first + count <= index_count());

			if (is_instanced())
			{
				assert(instance_count <= this->instance_count());

				if (count != 0 && instance_count != 0)
				{
					bind();
					gl::DrawElementsInstanced(primitive_type(), count, index_type(), reinterpret_cast<char const *>(detail::index_size(index_type()) * first), instance_count);
				}
			}
			else
			{
				if (count != 0)
				{
					bind();
					gl::DrawElements(primitive_type(), count, index_type(), reinterpret_cast<char const *>(detail::index_size(index_type()) * first));
				}
			}
		}
		else
		{
			assert(first + count <= vertex_count());

			if (is_instanced())
			{
				assert(instance_count <= this->instance_count());

				if (count != 0 && instance_count != 0)
				{
					bind();
					gl::DrawArraysInstanced(primitive_type(), first, count, instance_count);
				}
			}
			else
			{
				if (count != 0)
				{
					bind();
					gl::DrawArrays(primitive_type(), first, count);
				}
			}
		}
	}

	imported_mesh load_mesh(std::string_view data)
	{
		// These should be in sync with convert-mesh.py
		static std::uint32_t const SECTION_MESH  = 1;
		static std::uint32_t const SECTION_BONES = 2;
		static std::uint32_t const SECTION_POSE  = 3;

		static std::uint32_t const POSITION_MASK  = 1 << 0;
		static std::uint32_t const NORMALS_MASK   = 1 << 1;
		static std::uint32_t const COLORS_MASK    = 1 << 2;
		static std::uint32_t const TEXCOORDS_MASK = 1 << 3;
		static std::uint32_t const WEIGHTS_MASK   = 1 << 4;

		std::uint32_t vertex_format = 0;

		util::binary_istream s{data};

		imported_mesh result;

		auto parse_section_mesh = [&]
		{
			vertex_format = s.read<std::uint32_t>();

			if (vertex_format & POSITION_MASK)
				result.attribs += make_attribs_description<math::point<float, 3>>();
			else
				result.attribs += make_attribs_description<gfx::skip>();

			if (vertex_format & NORMALS_MASK)
				result.attribs += make_attribs_description<math::vector<float, 3>>();
			else
				result.attribs += make_attribs_description<gfx::skip>();

			if (vertex_format & COLORS_MASK)
				result.attribs += make_attribs_description<gfx::normalized<math::vector<std::uint8_t, 4>>>();
			else
				result.attribs += make_attribs_description<gfx::skip>();

			if (vertex_format & TEXCOORDS_MASK)
				result.attribs += make_attribs_description<math::vector<float, 2>>();
			else
				result.attribs += make_attribs_description<gfx::skip>();

			if (vertex_format & WEIGHTS_MASK)
				result.attribs += make_attribs_description<gfx::integer<math::vector<std::uint8_t, 2>>, gfx::normalized<math::vector<std::uint8_t, 2>>>();
			else
				result.attribs += make_attribs_description<gfx::skip, gfx::skip>();

			auto vertex_count = s.read<std::uint32_t>();
			auto vertex_ptr = s.read_raw(vertex_count * result.attribs.vertex_size);

			auto index_count = s.read<std::uint32_t>();
			auto index_ptr = s.read_ptr<std::uint32_t>(index_count);

			result.vertices = {vertex_ptr, vertex_ptr + vertex_count * result.attribs.vertex_size};
			result.indices = {index_ptr, index_ptr + index_count};
		};

		auto parse_section_bones = [&]
		{
			auto bone_count = s.read<std::uint32_t>();
			auto bone_ptr = s.read_ptr<bone>(bone_count);

			result.bones = {bone_ptr, bone_ptr + bone_count};
		};

		auto parse_section_pose = [&]
		{
			auto name_length = s.read<std::uint32_t>();
			auto name_ptr = s.read_raw(name_length);

			auto pose_count = s.read<std::uint32_t>();
			if (pose_count != result.bones.size())
				throw util::exception("Number of transforms in a pose must be equal to the number of bones");
			auto pose_ptr = s.read_ptr<bone_transform<float>>(pose_count);

			std::string_view name(name_ptr, name_length);

			result.poses[name] = {pose_ptr, pose_ptr + pose_count};
		};

		bool had_section_mesh = false;
		bool had_section_bones = false;

		while (!s.eof())
		{
			auto section_type = s.read<std::uint32_t>();

			switch (section_type)
			{
			case SECTION_MESH:
				if (had_section_mesh)
					throw util::exception("Section 'mesh' must not repeat");
				parse_section_mesh();
				had_section_mesh = true;
				break;
			case SECTION_BONES:
				if (had_section_bones)
					throw util::exception("Section 'bones' must not repeat");
				if (!had_section_mesh)
					throw util::exception("Section 'bones' must come after section 'mesh'");
				if ((vertex_format & WEIGHTS_MASK) == 0)
					throw util::exception("Section 'bones' requires weights in vertex format");
				parse_section_bones();
				had_section_bones = true;
				break;
			case SECTION_POSE:
				if (!had_section_bones)
					throw util::exception("Section 'pose' must come after section 'bones'");
				parse_section_pose();
				break;
			default:
				throw util::exception("Unknown section code " + util::to_string(section_type));
			}
		}

		if (!had_section_mesh)
			throw util::exception("Section 'mesh' must be present");

		return result;
	}

}