#include "gr2_bridge.h"

#include <godot_cpp/classes/array_mesh.hpp>
#include <godot_cpp/classes/mesh.hpp>
#include <godot_cpp/classes/skeleton3d.hpp>
#include <godot_cpp/classes/skin.hpp>
#include <godot_cpp/core/math.hpp>
#include <godot_cpp/core/memory.hpp>
#include <godot_cpp/variant/packed_float32_array.hpp>
#include <godot_cpp/variant/packed_int32_array.hpp>
#include <godot_cpp/variant/packed_vector2_array.hpp>
#include <godot_cpp/variant/packed_vector3_array.hpp>
#include <godot_cpp/variant/utility_functions.hpp>

#include <vector>

using namespace godot;

namespace mtgodot {

gr2::Mat4 mul4x3(const gr2::Mat4 &A, const gr2::Mat4 &B) {
	gr2::Mat4 R{};
	for (int i = 0; i < 3; ++i) {
		for (int k = 0; k < 3; ++k)
			R[i * 4 + k] = A[i * 4 + 0] * B[0 * 4 + k] + A[i * 4 + 1] * B[1 * 4 + k] +
					A[i * 4 + 2] * B[2 * 4 + k];
	}
	for (int k = 0; k < 3; ++k)
		R[12 + k] = A[12] * B[k] + A[13] * B[4 + k] + A[14] * B[8 + k] + B[12 + k];
	R[15] = 1.0f;
	return R;
}

Transform3D gr2_to_godot(const gr2::Mat4 &m) {
	// gr2 row-major, row-vector: basis columns are (m0,m1,m2),(m4,m5,m6),(m8,m9,m10);
	// translation is the 4th row (m12,m13,m14).
	Basis b(
			Vector3(m[0], m[1], m[2]),
			Vector3(m[4], m[5], m[6]),
			Vector3(m[8], m[9], m[10]));
	return Transform3D(b, Vector3(m[12], m[13], m[14]));
}

Transform3D make_conv(float unit_scale, bool flip_z) {
	// Z-up -> Y-up: rotate -90 deg about X. Then uniform scale. Optional Z flip.
	Basis b;
	b = b.rotated(Vector3(1, 0, 0), Math::deg_to_rad(-90.0));
	b = b.scaled(Vector3(unit_scale, unit_scale, flip_z ? -unit_scale : unit_scale));
	return Transform3D(b, Vector3());
}

Skeleton3D *build_skeleton(const gr2::Skeleton &sk) {
	if (sk.bones.empty()) {
		return nullptr;
	}
	Skeleton3D *skel = memnew(Skeleton3D);
	skel->set_name("Skeleton3D");

	const Transform3D ip = gr2_to_godot(sk.initial_placement);

	for (size_t i = 0; i < sk.bones.size(); ++i) {
		skel->add_bone(String(sk.bones[i].name.c_str()));
	}
	for (size_t i = 0; i < sk.bones.size(); ++i) {
		const gr2::Bone &bn = sk.bones[i];
		skel->set_bone_parent((int)i, bn.parent);
		Transform3D rest = gr2_to_godot(bn.local_transform);
		if (bn.parent < 0) {
			rest = ip * rest; // fold initial_placement into root bones
		}
		skel->set_bone_rest((int)i, rest);
		skel->reset_bone_pose((int)i);
	}
	return skel;
}

Ref<Skin> build_skin(const gr2::Skeleton &sk) {
	Ref<Skin> skin;
	skin.instantiate();
	for (size_t i = 0; i < sk.bones.size(); ++i) {
		skin->add_bind((int)i, gr2_to_godot(sk.bones[i].inverse_world));
		skin->set_bind_name((int)i, StringName(sk.bones[i].name.c_str()));
	}
	return skin;
}

std::vector<RenderPart> build_parts(const gr2::FileInfo &fi) {
	std::vector<RenderPart> parts;
	for (int mi = 0; mi < (int)fi.meshes.size(); ++mi) {
		const gr2::Mesh &m = fi.meshes[mi];
		if (m.vertices.empty() || m.indices.empty()) {
			continue;
		}
		const uint32_t total = (uint32_t)m.indices.size();
		if (m.tri_groups.size() <= 1) {
			RenderPart p;
			p.mesh = mi;
			p.mat_index = m.tri_groups.empty() ? -1 : m.tri_groups[0].material_index;
			p.idx_first = 0;
			p.idx_count = total;
			parts.push_back(p);
			continue;
		}
		for (int g = 0; g < (int)m.tri_groups.size(); ++g) {
			const gr2::TriGroup &tg = m.tri_groups[g];
			if (tg.tri_count <= 0) {
				continue;
			}
			uint32_t first = (uint32_t)(tg.tri_first < 0 ? 0 : tg.tri_first) * 3u;
			uint32_t count = (uint32_t)tg.tri_count * 3u;
			if (first >= total) {
				continue;
			}
			if (first + count > total) {
				count = total - first;
			}
			RenderPart p;
			p.mesh = mi;
			p.group = g;
			p.mat_index = tg.material_index;
			p.idx_first = first;
			p.idx_count = count;
			parts.push_back(p);
		}
	}
	return parts;
}

Ref<ArrayMesh> build_mesh(const gr2::FileInfo &fi, const std::vector<RenderPart> &parts,
		bool flip_winding, AABB &out_bounds) {
	Ref<ArrayMesh> am;
	am.instantiate();
	am->set_name("ArrayMesh");

	bool have_bounds = false;

	// vertex arrays are per gr2 mesh; multiple parts of one mesh reuse them.
	for (const RenderPart &part : parts) {
		const gr2::Mesh &m = fi.meshes[part.mesh];
		const int vcount = (int)m.vertices.size();

		PackedVector3Array pos;
		PackedVector3Array nrm;
		PackedVector2Array uv;
		PackedInt32Array bones;
		PackedFloat32Array weights;
		pos.resize(vcount);
		nrm.resize(vcount);
		uv.resize(vcount);
		bones.resize(vcount * 4);
		weights.resize(vcount * 4);

		Vector3 *pos_w = pos.ptrw();
		Vector3 *nrm_w = nrm.ptrw();
		Vector2 *uv_w = uv.ptrw();
		int32_t *bn_w = bones.ptrw();
		float *wt_w = weights.ptrw();

		auto slot_to_bone = [&](int slot) -> int {
			if (slot < 0 || slot >= (int)m.bone_bindings.size()) {
				return 0;
			}
			int b = m.bone_bindings[slot];
			return (b < 0) ? 0 : b;
		};

		const int rigid_bone = m.rigid ? slot_to_bone(m.bone_bindings.empty() ? -1 : 0) : 0;

		for (int i = 0; i < vcount; ++i) {
			const gr2::Vertex &v = m.vertices[i];
			pos_w[i] = Vector3(v.pos[0], v.pos[1], v.pos[2]);
			nrm_w[i] = Vector3(v.normal[0], v.normal[1], v.normal[2]);
			uv_w[i] = Vector2(v.uv0[0], v.uv0[1]);

			if (!have_bounds) {
				out_bounds.position = pos_w[i];
				out_bounds.size = Vector3();
				have_bounds = true;
			} else {
				out_bounds = out_bounds.expand(pos_w[i]);
			}

			if (m.rigid) {
				bn_w[i * 4 + 0] = rigid_bone;
				bn_w[i * 4 + 1] = bn_w[i * 4 + 2] = bn_w[i * 4 + 3] = 0;
				wt_w[i * 4 + 0] = 1.0f;
				wt_w[i * 4 + 1] = wt_w[i * 4 + 2] = wt_w[i * 4 + 3] = 0.0f;
				continue;
			}

			float wsum = 0.0f;
			for (int k = 0; k < 4; ++k) {
				wsum += v.bone_weight[k];
			}
			for (int k = 0; k < 4; ++k) {
				bn_w[i * 4 + k] = slot_to_bone(v.bone_index[k]);
				wt_w[i * 4 + k] = (wsum > 0.0f) ? (v.bone_weight[k] / (float)wsum) : (k == 0 ? 1.0f : 0.0f);
			}
		}

		// indices: this part's sub-range of mesh.indices only
		const uint32_t ib = part.idx_first;
		const uint32_t ic = (part.idx_count && part.idx_first + part.idx_count <= m.indices.size())
				? part.idx_count
				: (uint32_t)m.indices.size() - ib;
		PackedInt32Array idx;
		idx.resize((int)ic);
		int32_t *idx_w = idx.ptrw();
		if (flip_winding) {
			for (uint32_t t = 0; t + 2 < ic; t += 3) {
				idx_w[t + 0] = (int)m.indices[ib + t + 0];
				idx_w[t + 1] = (int)m.indices[ib + t + 2];
				idx_w[t + 2] = (int)m.indices[ib + t + 1];
			}
		} else {
			for (uint32_t t = 0; t < ic; ++t) {
				idx_w[t] = (int)m.indices[ib + t];
			}
		}

		Array arrays;
		arrays.resize(Mesh::ARRAY_MAX);
		arrays[Mesh::ARRAY_VERTEX] = pos;
		arrays[Mesh::ARRAY_NORMAL] = nrm;
		arrays[Mesh::ARRAY_TEX_UV] = uv;
		arrays[Mesh::ARRAY_BONES] = bones;
		arrays[Mesh::ARRAY_WEIGHTS] = weights;
		arrays[Mesh::ARRAY_INDEX] = idx;

		am->add_surface_from_arrays(Mesh::PRIMITIVE_TRIANGLES, arrays);
		String nm = m.name.empty() ? String("surf_") + itos(part.mesh) : String(m.name.c_str());
		if (part.group >= 0) {
			nm += String("#") + itos(part.group);
		}
		am->surface_set_name(am->get_surface_count() - 1, nm);
	}

	if (!have_bounds) {
		out_bounds = AABB();
	}
	return am;
}

} // namespace mtgodot