#include <psemek/gfx/gltf_parser.hpp>
#include <psemek/math/scale.hpp>
#include <psemek/math/rotation.hpp>
#include <psemek/math/translation.hpp>
#include <psemek/util/to_string.hpp>
#include <psemek/util/hash_table.hpp>
#include <psemek/util/exception.hpp>
#include <psemek/io/memory_stream.hpp>
#include <psemek/log/log.hpp>

#include <boost/preprocessor/stringize.hpp>

#define RAPIDJSON_ASSERT(x) if (!(x)) throw ::psemek::util::exception("Error parsing glTF: " BOOST_PP_STRINGIZE(x));
#define RAPIDJSON_NOEXCEPT_ASSERT(x)

#include <rapidjson/document.h>
#include <rapidjson/istreamwrapper.h>

#include <unordered_set>

namespace psemek::gfx
{

	namespace
	{

		std::string_view to_string(rapidjson::ParseErrorCode const & code)
		{
			using namespace rapidjson;

			switch (code)
			{
			case kParseErrorNone: return "no error";
			case kParseErrorDocumentEmpty: return "document is empty";
			case kParseErrorDocumentRootNotSingular: return "document root must not be followed by other values";
			case kParseErrorValueInvalid: return "invalid value";
			case kParseErrorObjectMissName: return "missing a name for object member";
			case kParseErrorObjectMissColon: return "missing a colon after a name of object member";
			case kParseErrorObjectMissCommaOrCurlyBracket: return "missing a comma or '}' after an object member";
			case kParseErrorArrayMissCommaOrSquareBracket: return "missing a comma or ']' after an array element";
			case kParseErrorStringUnicodeEscapeInvalidHex: return R"(Incorrect hex digit after \\u escape in string)";
			case kParseErrorStringUnicodeSurrogateInvalid: return "surrogate pair in string is invalid";
			case kParseErrorStringEscapeInvalid: return "invalid escape character in string";
			case kParseErrorStringMissQuotationMark: return "missing a closing quotation mark in string";
			case kParseErrorStringInvalidEncoding: return "invalid encoding in string";
			case kParseErrorNumberTooBig: return "number too big to be stored in double";
			case kParseErrorNumberMissFraction: return "missing fraction part in number";
			case kParseErrorNumberMissExponent: return "missing exponent in number";
			case kParseErrorTermination: return "parsing was terminated";
			case kParseErrorUnspecificSyntaxError: return "unspecific syntax error";
			default: return "(unknown error)";
			}
		}

		auto parse_accessor_type(std::string const & type)
		{
			if (type == "SCALAR")
				return gltf_asset::accessor::scalar;
			if (type == "VEC2")
				return gltf_asset::accessor::vec2;
			if (type == "VEC3")
				return gltf_asset::accessor::vec3;
			if (type == "VEC4")
				return gltf_asset::accessor::vec4;
			if (type == "MAT2")
				return gltf_asset::accessor::mat2;
			if (type == "MAT3")
				return gltf_asset::accessor::mat3;
			if (type == "MAT4")
				return gltf_asset::accessor::mat4;
			throw util::exception(util::to_string("Unknown accessor component type: ", type));
		}

		static util::hash_set<std::string> const supported_extensions =
		{
			"KHR_lights_punctual",
			"KHR_materials_emissive_strength"
		};

		gltf_asset::extras_map parse_extras(rapidjson::GenericValue<rapidjson::UTF8<>> const & value)
		{
			gltf_asset::extras_map result;

			if (value.HasMember("extras"))
			{
				for (auto const & extra : value["extras"].GetObject())
				{
					std::string name = extra.name.GetString();
					gltf_asset::extra value;

					bool error = false;

					if (extra.value.IsBool())
						value = extra.value.GetBool();
					else if (extra.value.IsNumber())
						value = extra.value.GetFloat();
					else if (extra.value.IsString())
						value = extra.value.GetString();
					else if (extra.value.IsArray())
					{
						std::vector<float> v;
						for (auto const & val : extra.value.GetArray())
						{
							if (!val.IsNumber())
								error = true;
							else
								v.push_back(val.GetFloat());
						}
						value = std::move(v);
					}
					else
						error = true;

					if (error)
						log::warning() << "every 'extras' value must be either a boolean, a number, an array of numbers, or a string (while parsing glTF)";
					else
						result[name] = std::move(value);
				}
			}

			return result;
		}

	}

	gltf_asset parse_gltf(io::istream && stream)
	{
		auto description_str = io::read_full(std::move(stream)).string();

		rapidjson::Document document;
		document.ParseInsitu(description_str.data());

		if (document.HasParseError())
			throw util::exception(util::to_string("Error parsing glTF: ", to_string(document.GetParseError()), " at ", document.GetErrorOffset()));

		gltf_asset result;

		if (document.HasMember("extensionsRequired"))
			for (auto const & extension : document["extensionsRequired"].GetArray())
				if (!supported_extensions.contains(extension.GetString()))
					throw util::exception("glTF extension " + std::string(extension.GetString()) + " is not supported");

		if (document.HasMember("nodes")) for (auto const & node : document["nodes"].GetArray())
		{
			auto & target = result.nodes.emplace_back();
			target.name = node["name"].GetString();
			if (node.HasMember("mesh"))
				target.mesh = node["mesh"].GetUint64();

			if (node.HasMember("skin"))
				target.skin = node["skin"].GetUint64();

			target.translation = {0.f, 0.f, 0.f};
			if (node.HasMember("translation"))
			{
				auto const & translation = node["translation"].GetArray();
				for (std::size_t i = 0; i < 3; ++i)
					target.translation[i] = translation[i].GetFloat();
			}

			target.rotation = math::quaternion<float>::identity();
			if (node.HasMember("rotation"))
			{
				auto const & rotation = node["rotation"].GetArray();
				for (std::size_t i = 0; i < 4; ++i)
					target.rotation[i] = rotation[i].GetFloat();
			}

			target.scale = {1.f, 1.f, 1.f};
			if (node.HasMember("scale"))
			{
				auto const & scale = node["scale"].GetArray();
				for (std::size_t i = 0; i < 3; ++i)
					target.scale[i] = scale[i].GetFloat();
			}

			target.transform = math::translation(target.translation).transform()
				* math::quaternion_rotation(target.rotation).transform()
				* math::scale(target.scale).transform();

			if (node.HasMember("children"))
			{
				for (auto const & child : node["children"].GetArray())
					target.children.push_back(child.GetUint64());
			}

			if (node.HasMember("extensions"))
				for (auto const & extension : node["extensions"].GetObject())
				{
					if (extension.name == "KHR_lights_punctual")
					{
						target.light = extension.value["light"].GetUint64();
					}
				}

			target.extras = parse_extras(node);
		}

		if (document.HasMember("meshes")) for (auto const & mesh : document["meshes"].GetArray())
		{
			auto & target = result.meshes.emplace_back();
			for (auto const & primitive : mesh["primitives"].GetArray())
			{
				auto & primitive_target = target.primitives.emplace_back();
				primitive_target.indices = primitive["indices"].GetUint64();
				if (primitive.HasMember("material"))
					primitive_target.material = primitive["material"].GetUint64();

				auto const & attributes = primitive["attributes"];
				if (attributes.HasMember("POSITION"))
					primitive_target.position = attributes["POSITION"].GetUint64();
				if (attributes.HasMember("NORMAL"))
					primitive_target.normal = attributes["NORMAL"].GetUint64();
				if (attributes.HasMember("TEXCOORD_0"))
					primitive_target.texcoord = attributes["TEXCOORD_0"].GetUint64();
				if (attributes.HasMember("COLOR_0"))
					primitive_target.color = attributes["COLOR_0"].GetUint64();
				if (attributes.HasMember("JOINTS_0"))
					primitive_target.joints = attributes["JOINTS_0"].GetUint64();
				if (attributes.HasMember("WEIGHTS_0"))
					primitive_target.weights = attributes["WEIGHTS_0"].GetUint64();
			}
		}

		if (document.HasMember("materials")) for (auto const & material : document["materials"].GetArray())
		{
			auto & target = result.materials.emplace_back();

			target.name = material["name"].GetString();

			if (material.HasMember("doubleSided"))
				target.two_sided = material["doubleSided"].GetBool();
			else
				target.two_sided = false;

			auto const & pbr = material["pbrMetallicRoughness"];

			if (pbr.HasMember("baseColorFactor"))
			{
				auto const & color = pbr["baseColorFactor"].GetArray();
				gfx::color_4f & target_color = target.albedo.emplace();
				for (std::size_t i = 0; i < 4; ++i)
					target_color[i] = color[i].GetFloat();
			}

			if (pbr.HasMember("baseColorTexture"))
			{
				target.texture = document["textures"].GetArray()[pbr["baseColorTexture"]["index"].GetUint64()]["source"].GetInt64();
			}

			if (pbr.HasMember("metallicFactor"))
			{
				target.metallic = pbr["metallicFactor"].GetFloat();
			}

			if (pbr.HasMember("roughnessFactor"))
			{
				target.roughness = pbr["roughnessFactor"].GetFloat();
			}

			if (pbr.HasMember("metallicRoughnessTexture"))
			{
				target.material_texture = document["textures"].GetArray()[pbr["metallicRoughnessTexture"]["index"].GetUint64()]["source"].GetInt64();
			}

			if (material.HasMember("emissiveFactor"))
			{
				auto const & emission = material["emissiveFactor"].GetArray();
				gfx::color_3f & target_emission = target.emission.emplace();
				for (std::size_t i = 0; i < 3; ++i)
					target_emission[i] = emission[i].GetFloat();
			}

			if (material.HasMember("emissiveTexture"))
			{
				target.emission_texture = document["textures"].GetArray()[material["emissiveTexture"]["index"].GetUint64()]["source"].GetInt64();
			}

			if (material.HasMember("extensions"))
			{
				auto const & extensions = material["extensions"];
				if (extensions.HasMember("KHR_materials_emissive_strength"))
				{
					if (!target.emission)
						target.emission = {1.f, 1.f, 1.f};
					*target.emission *= extensions["KHR_materials_emissive_strength"]["emissiveStrength"].GetFloat();
				}
			}

			target.extras = parse_extras(material);
		}

		if (document.HasMember("images")) for (auto const & image : document["images"].GetArray())
		{
			auto & target = result.textures.emplace_back();
			if (image.HasMember("uri"))
				target.uri = image["uri"].GetString();
			if (image.HasMember("bufferView"))
				target.buffer_view = image["bufferView"].GetUint64();
		}

		if (document.HasMember("skins")) for (auto const & skin : document["skins"].GetArray())
		{
			auto & target = result.skins.emplace_back();

			for (auto const & joint : skin["joints"].GetArray())
				target.joints.push_back(joint.GetUint64());

			if (skin.HasMember("inverseBindMatrices"))
				target.inverse_bind_matrices = skin["inverseBindMatrices"].GetUint64();

			if (skin.HasMember("name"))
				target.name = skin["name"].GetString();
		}

		if (document.HasMember("animations")) for (auto const & animation : document["animations"].GetArray())
		{
			auto & target = result.animations.emplace_back();

			if (animation.HasMember("name"))
				target.name = animation["name"].GetString();

			for (auto const & channel : animation["channels"].GetArray())
			{
				auto & channel_target = target.channels.emplace_back();

				channel_target.sampler = channel["sampler"].GetUint64();
				channel_target.target = channel["target"]["node"].GetUint64();
				std::string path{channel["target"]["path"].GetString()};
				if (path == "scale")
					channel_target.path = gltf_asset::animation::channel::scale;
				else if (path == "rotation")
					channel_target.path = gltf_asset::animation::channel::rotation;
				else if (path == "translation")
					channel_target.path = gltf_asset::animation::channel::translation;
				else
					throw util::exception("Unknown glTF animation target path: " + path);
			}

			for (auto const & sampler : animation["samplers"].GetArray())
			{
				auto & sampler_target = target.samplers.emplace_back();

				sampler_target.input = sampler["input"].GetUint64();
				sampler_target.output = sampler["output"].GetUint64();

				std::string interpolation{sampler["interpolation"].GetString()};
				if (interpolation == "STEP")
					sampler_target.interpolation = math::easing_type::constant_left;
				else if (interpolation == "LINEAR")
					sampler_target.interpolation = math::easing_type::linear;
				else if (interpolation == "CUBICSPLINE")
					sampler_target.interpolation = math::easing_type::cubic;
				else
					throw util::exception("Unknown glTF animation interpolation type: " + interpolation);
			}
		}

		if (document.HasMember("accessors")) for (auto const & accessor : document["accessors"].GetArray())
		{
			auto & target = result.accessors.emplace_back();

			target.buffer_view = accessor["bufferView"].GetUint64();
			target.component_type = (gltf_asset::accessor::component_type_t)accessor["componentType"].GetUint64();
			target.count = accessor["count"].GetUint64();
			target.type = parse_accessor_type(accessor["type"].GetString());

			if (accessor.HasMember("normalized"))
				target.normalized = accessor["normalized"].GetBool();
			else
				target.normalized = false;
		}

		if (document.HasMember("bufferViews")) for (auto const & buffer_view : document["bufferViews"].GetArray())
		{
			auto & target = result.buffer_views.emplace_back();

			target.buffer = buffer_view["buffer"].GetUint64();
			target.offset = buffer_view["byteOffset"].GetUint64();
			target.length = buffer_view["byteLength"].GetUint64();
			target.stride = 0;
			if (buffer_view.HasMember("byteStride"))
				target.stride = buffer_view["byteStride"].GetUint64();
		}

		if (document.HasMember("buffers")) for (auto const & buffer : document["buffers"].GetArray())
		{
			auto & target = result.buffers.emplace_back();

			target.length = buffer["byteLength"].GetUint64();

			if (buffer.HasMember("uri"))
				target.uri = buffer["uri"].GetString();
		}

		if (document.HasMember("extensions")) for (auto const & extension : document["extensions"].GetObject())
		{
			if (extension.name == "KHR_lights_punctual")
			{
				for (auto const & light : extension.value["lights"].GetArray())
				{
					auto & target = result.lights.emplace_back();
					target.color = {1.f, 1.f, 1.f};
					target.intensity = 1.f;
					target.range = std::numeric_limits<float>::infinity();
					target.cone_angle = {0.f, math::pi / 4.f};

					std::string type = light["type"].GetString();

					if (type == "directional")
						target.type = gltf_asset::light::directional;
					else if (type == "point")
						target.type = gltf_asset::light::point;
					else if (type == "spot")
						target.type = gltf_asset::light::spot;
					else
						throw util::exception("Unknown light type: " + type);

					if (light.HasMember("color"))
					{
						auto const & color = light["color"].GetArray();
						for (std::size_t i = 0; i < 3; ++i)
							target.color[i] = color[i].GetFloat();
					}

					if (light.HasMember("intensity"))
						target.intensity = light["intensity"].GetFloat();

					if (light.HasMember("range"))
						target.range = light["range"].GetFloat();

					if (target.type == gltf_asset::light::spot)
					{
						auto const & spot = light["spot"].GetObject();

						if (spot.HasMember("innerConeAngle"))
							target.cone_angle.min = spot["innerConeAngle"].GetFloat();

						if (spot.HasMember("outerConeAngle"))
							target.cone_angle.max = spot["outerConeAngle"].GetFloat();
					}

					target.extras = parse_extras(light);
				}
			}
		}

		for (std::size_t i = 0; i < result.nodes.size(); ++i)
			for (auto child : result.nodes[i].children)
				result.nodes[child].parent = i;

		for (std::size_t i = 0; i < result.nodes.size(); ++i)
			result.node_index[result.nodes[i].name] = i;

		for (std::size_t i = 0; i < result.materials.size(); ++i)
			result.material_index[result.materials[i].name] = i;

		result.node_animations.resize(result.nodes.size());
		for (std::size_t i = 0; i < result.animations.size(); ++i)
			for (std::size_t j = 0; j < result.animations[i].channels.size(); ++j)
				result.node_animations[result.animations[i].channels[j].target].push_back({i, j});

		return result;
	}

	gltf_asset parse_glb(io::istream && stream)
	{
		std::uint32_t magic, version, length;
		stream.read_all((char *)&magic, 4);
		stream.read_all((char *)&version, 4);
		stream.read_all((char *)&length, 4);

		if (magic != 0x46546c67u)
			throw util::exception("Error parsing GLB magic");

		if (version != 2)
			throw util::exception("Unknown GLB file version");

		length -= 12;

		gltf_asset result;

		while (length > 0)
		{
			std::uint32_t chunk_length, chunk_type;
			stream.read_all((char *)&chunk_length, 4);
			stream.read_all((char *)&chunk_type, 4);

			util::blob data(chunk_length);
			stream.read_all(data.data(), chunk_length);

			length -= 8;
			length -= chunk_length;

			if (chunk_type == 0x4e4f534au)
			{
				result = parse_gltf(io::memory_istream(data.begin(), data.end()));
			}
			else if (chunk_type == 0x004e4942u)
			{
				if (result.buffers.empty())
					throw util::exception("Error parsing GLB: glTF chunk has no buffers");

				if (!result.buffers[0].uri.empty())
					throw util::exception("Error parsing GLB: first glTF buffer has URI");

				result.buffers[0].data = std::move(data);
			}
			else
				throw util::exception("Error parsing GLB: unknown chunk type");
		}

		return result;
	}

}
