#!/usr/bin/env python3

import pygltflib
import sys
import numpy
import os
import base64
import pathlib
import struct

def mat3_to_quat(m):
	trace = numpy.trace(m)
	
	if trace >= 0:
		r = (1 + trace) ** 0.5
		rinv = 0.5 / r
		
		q0 = rinv * (m[1, 2] - m[2, 1])
		q1 = rinv * (m[2, 0] - m[0, 2])
		q2 = rinv * (m[0, 1] - m[1, 0])
		q3 = r * 0.5
	elif m[0, 0] >= m[1, 1] and m[0, 0] >= m[2, 2]:
		r = (1 - m[1, 1] - m[2, 2] + m[0, 0]) ** 0.5
		rinv = 0.5 / r
		
		q0 = r * 0.5
		q1 = rinv * (m[0, 1] + m[1, 0])
		q2 = rinv * (m[0, 2] + m[2, 0])
		q3 = rinv * (m[1, 2] - m[2, 1])
	elif m[1, 1] >= m[2, 2]:
		r = (1 - m[0, 0] - m[2, 2] + m[1, 1]) ** 0.5
		rinv = 0.5 / r
		
		q0 = rinv * (m[0, 1] + m[1, 0])
		q1 = r * 0.5
		q2 = rinv * (m[1, 2] + m[2, 1])
		q3 = rinv * (m[2, 0] - m[0, 2])
	else:
		r = (1 - m[0, 0] - m[1, 1] + m[2, 2]) ** 0.5
		rinv = 0.5 / r
		
		q0 = rinv * (m[0, 2] + m[2, 0])
		q1 = rinv * (m[1, 2] + m[2, 1])
		q2 = r * 0.5
		q3 = rinv * (m[0, 1] - m[1, 0])
	
	q = numpy.array([q0, q1, q2, q3])
	
	if q[3] < 0:
		q = -q
	
	return q

def mat3_from_quat(q):
	b, c, d, a = q[0], q[1], q[2], q[3]
	
	return numpy.array([
		[a ** 2 + b ** 2 - c ** 2 - d ** 2, 2 * b * c + 2 * a * d, 2 * b * d - 2 * a * c],
		[2 * b * c - 2 * a * d, a ** 2 - b ** 2 + c ** 2 - d ** 2, 2 * c * d - 2 * a * b],
		[2 * b * d - 2 * a * c, 2 * c * d + 2 * a * b, a ** 2 - b ** 2 - c ** 2 + d ** 2],
	])

def mat4_decompose(m):
	translation = m[3][:3]
	
	rs_matrix = m[:3, :3]
	
	scale = (rs_matrix ** 2).sum(-1) ** 0.5
	
	rs_matrix[0] /= scale[0]
	rs_matrix[1] /= scale[1]
	rs_matrix[2] /= scale[2]
	
	if numpy.dot(numpy.cross(m[0][:3], m[1][:3]), m[2][:3]) < 0:
		rs_matrix[0] *= -1
		rs_matrix[1] *= -1
		rs_matrix[2] *= -1
		scale *= -1
	
	return translation, mat3_to_quat(rs_matrix), scale

def mat4_compose(t, r, s):
	m = numpy.zeros((4, 4))
	
	rm = mat3_from_quat(r) if r is not None else numpy.eye(3)
	
	if s is not None:
		rm[0] *= s[0]
		rm[1] *= s[1]
		rm[2] *= s[2]
	
	m[:3, :3] = rm
	
	if t is not None:
		m[3][:3] = t
	
	m[3][3] = 1
	
	return m

def lerp(a, b, alpha):
	return a + (b - a) * alpha

def slerp(a, b, alpha):
	angle = numpy.arccos(min(1, max(-1, numpy.dot(a, b))))
	denominator = numpy.sin(angle)
	if abs(denominator) < 0.000001:
		return lerp(a, b, alpha)
	return (a * numpy.sin((1 - alpha) * angle) + b * numpy.sin(alpha * angle)) / denominator

TYPE_SIZES = {
	"SCALAR": 1,
	"VEC2": 2,
	"VEC3": 3,
	"VEC4": 4,
	"MAT2": 4,
	"MAT4": 16,
	"MAT3": 9,
	pygltflib.FLOAT: 4,
	pygltflib.UNSIGNED_INT: 4,
	pygltflib.UNSIGNED_SHORT: 2,
	pygltflib.UNSIGNED_BYTE: 1,
	pygltflib.SHORT: 2,
	pygltflib.BYTE: 1,
}

GLTF_NUMPY_MAPPING = {
	pygltflib.FLOAT: numpy.float32,
	pygltflib.UNSIGNED_INT: numpy.uint32,
	pygltflib.UNSIGNED_SHORT: numpy.uint16,
	pygltflib.UNSIGNED_BYTE: numpy.uint8,
	pygltflib.SHORT: numpy.int16,
	pygltflib.BYTE: numpy.int8,
}

errors = []
def err(s):
	errors.append(s)
def err_barrier():
	for e in errors:
		print(e)
	if os.name == "nt" and "PROMPT" not in os.environ:
		raw_input()
	if len(errors):
		sys.exit(1)

for input_filename in sys.argv[1:]:
	gltf = pygltflib.GLTF2().load(input_filename)
	gltf.convert_buffers(pygltflib.BufferFormat.DATAURI)
	
	def get_nd(acc_idx):
		if acc_idx == None:
			return None
		
		accessor = gltf.accessors[acc_idx]
		buf_view = gltf.bufferViews[accessor.bufferView]
		buf = gltf.buffers[buf_view.buffer]
		uri = buf.uri
		
		if "base64," in uri:
			buf_data = base64.decodebytes(buf.uri.split("base64,")[-1].encode())
		else:
			with open(uri, "rb") as urifile:
				buf_data = urifile.read()
		
		start = buf_view.byteOffset + accessor.byteOffset
		array = numpy.frombuffer(buf_data[start : start + accessor.count * TYPE_SIZES[accessor.type] * TYPE_SIZES[accessor.componentType]], dtype = GLTF_NUMPY_MAPPING[accessor.componentType])
		return numpy.reshape(array, (array.shape[0] // TYPE_SIZES[accessor.type], TYPE_SIZES[accessor.type]))

	def get(acc_idx):
		array = get_nd(acc_idx)
		return array.ravel() if array is not None else None

	mesh_nodes = [n for n in gltf.nodes if n.mesh != None]
	
	if len(mesh_nodes) > 1:
		err("All meshes except first are ignored")
	
	mesh_node = mesh_nodes[0]

	mesh = gltf.meshes[mesh_node.mesh]
	skin = gltf.skins[mesh_node.skin] if mesh_node.skin != None else None
	
	k3result = {
		"animations": [],
		"bones": [],
		"meshes": [],
	}

	for primitive in mesh.primitives:
		material_name = ""
		
		if primitive.material is not None:
			material_name = gltf.materials[primitive.material].name
			if not material_name:
				err("Missing material name")
				material_name = ""
		else:
			err("Missing material information")
		
		mehs = {
			"positions": get_nd(primitive.attributes.POSITION),
			"normals": get_nd(primitive.attributes.NORMAL),
			"uvs": get_nd(primitive.attributes.TEXCOORD_0),
			"colors": get_nd(primitive.attributes.COLOR_0),
			"boneids": get_nd(primitive.attributes.JOINTS_0),
			"boneweights": get_nd(primitive.attributes.WEIGHTS_0),
			"indices": get(primitive.indices),
			"material_name": material_name
		}
		
		if mehs["normals"] is None:
			err("Missing normal data")
		if mehs["uvs"] is None:
			err("Missing UVs")
		
		k3result["meshes"].append(mehs)
	
	if len(skin.joints) > 255:
		err("Bone maximum is 255")
	
	err_barrier()
	
	joint_gltf_to_k3_id_mapping = {}
	joint_k3_to_gltf_id_mapping = {}
	if skin:
		# Topologically sort joints just in case they aren't
		joint_parents = {i: next((j for j in skin.joints if i in gltf.nodes[j].children), None) for i in skin.joints}
		joint_graph = joint_parents.copy()
		while len(k3result["bones"]) != len(skin.joints):
			minimal_elements = [k for k, v in joint_graph.items() if v == None]
			for j in minimal_elements:
				joint_gltf_to_k3_id_mapping[j] = len(k3result["bones"])
				joint_k3_to_gltf_id_mapping[len(k3result["bones"])] = j
				k3result["bones"].append({
					"parent": joint_gltf_to_k3_id_mapping.get(joint_parents[j], 255),
					"bind": mat4_compose(gltf.nodes[j].translation, gltf.nodes[j].rotation, gltf.nodes[j].scale)
				})
			for k, v in joint_graph.copy().items():
				if v in minimal_elements:
					joint_graph[k] = None
				if v == None:
					del joint_graph[k]

	for k3b in k3result["bones"]:
		if k3b["parent"] != 255:
			k3b["bind"] = k3result["bones"][k3b["parent"]]["bind"] @ k3b["bind"]

	for anim_idx, animation in enumerate(gltf.animations):
		timelines = {}
		duration = 0
		
		# Fill timelines from GLTF data
		
		for channel in animation.channels:
			target_node = channel.target.node
			
			if target_node not in joint_gltf_to_k3_id_mapping:
				# Not a bone
				continue
			
			sampler = animation.samplers[channel.sampler]
			times = get(sampler.input)
			values = get_nd(sampler.output)
			
			timeline = []
			
			for t, v in zip(times, values):
				timeline.append((t, v))
				duration = max(duration, t)
			
			timeline.sort(key = lambda kf: kf[0])
			
			timelines.setdefault(joint_gltf_to_k3_id_mapping[target_node], {})[channel.target.path] = timeline
		
		# Make sure that all bones & transform types have at least one keyframe
		
		for bone_id, k3b in enumerate(k3result["bones"]):
			bone_timelines = timelines.setdefault(bone_id, {})
			l = bone_timelines.setdefault("translation", [])
			if len(l) == 0:
				l.append((0, numpy.array([0, 0, 0])))
			l = bone_timelines.setdefault("rotation", [])
			if len(l) == 0:
				l.append((0, numpy.array([0, 0, 0, 1])))
			l = bone_timelines.setdefault("scale", [])
			if len(l) == 0:
				l.append((0, numpy.array([1, 1, 1])))
		
		# Playback and store all bone transforms per frame as k3 format needs
		
		fps = 30
		frame_interval = 1 / fps
		
		frame_times = numpy.arange(0, duration, frame_interval)
		
		k3anim = {"id": anim_idx, "fps": fps, "frames": []}
		
		for t in frame_times:
			bone_transfs = []
			
			for bone_id, k3b in enumerate(k3result["bones"]):
				def dothing(transf_type, interp_func):
					tm = timelines[bone_id][transf_type]
					k1 = [k for k in range(len(tm)) if tm[k][0] <= t]
					k2 = [k for k in range(len(tm)) if tm[k][0] > t]
					k1 = max(k1) if len(k1) else 0
					k2 = min(k2) if len(k2) else len(tm) - 1
					alpha = (t - tm[k1][0]) / (tm[k2][0] - tm[k1][0]) if tm[k2][0] != tm[k1][0] else 0
					return interp_func(tm[k1][1], tm[k2][1], alpha)
				translation = dothing("translation", lerp)
				rotation = dothing("rotation", slerp)
				scale = dothing("scale", lerp)
				
				bone_transfs.append((translation, rotation))
			
			k3anim["frames"].append(bone_transfs)
			
		k3result["animations"].append(k3anim)

	# Finally, output
	
	output_filename = pathlib.Path(input_filename).with_suffix(".k3m")
	
	with open(output_filename, "wb") as f:
		f.write(b"K3M ")
		
		vertex_count = sum([len(mesh["positions"]) for mesh in k3result["meshes"]])
		index_count = sum([len(mesh["indices"]) for mesh in k3result["meshes"]])
		
		f.write(struct.pack("I", vertex_count))
		f.write(struct.pack("I", index_count))
		f.write(struct.pack("B", len(k3result["bones"])))
		
		colorsEnabled = any([mesh["colors"] is not None for mesh in k3result["meshes"]])
		
		# todo
		colorsEnabled = False
		
		f.write(struct.pack("B", 1 if colorsEnabled else 0))
		
		f.write(struct.pack("H", len(k3result["animations"])))
		
		for k3b in k3result["bones"]:
			f.write(struct.pack("16f", *numpy.reshape(numpy.linalg.inv(k3b["bind"]), (16,))))
		for k3b in k3result["bones"]:
			f.write(struct.pack("B", k3b["parent"]))
		
		for mesh in k3result["meshes"]:
			f.write(mesh["positions"][:, :3].astype(numpy.float32).ravel().tobytes())
		
		for mesh in k3result["meshes"]:
			normals = mesh["normals"][:, :3].astype(numpy.float32)
			for i in range(len(normals)):
				normals[i] /= numpy.max(numpy.abs(normals[i]))
				normals[i] *= 127
			f.write(normals.astype(numpy.int8).ravel().tobytes())
		
		for mesh in k3result["meshes"]:
			f.write(mesh["uvs"][:, :2].astype(numpy.float32).ravel().tobytes())
		
		if mesh["boneids"] is not None:
			for mesh in k3result["meshes"]:
				f.write(mesh["boneids"][:, :4].astype(numpy.int8).ravel().tobytes())
			
			for mesh in k3result["meshes"]:
				f.write((mesh["boneweights"][:, :4].astype(numpy.float32) * 65535).astype(numpy.int16).ravel().tobytes())
		
		if colorsEnabled:
			# TODO
			#for mesh in k3result["colors"]:
			#	f.write(mesh["colors"][:, :4])
			pass
		
		for mesh in k3result["meshes"]:
			f.write(mesh["indices"].astype(numpy.uint16).tobytes())
		
		f.write(struct.pack("H", len(k3result["meshes"])))
		
		offset = 0
		for mesh in k3result["meshes"]:
			ind_count = len(mesh["indices"])
			f.write(struct.pack("2H", offset, ind_count))
			f.write(mesh["material_name"].encode("UTF-8") + b'\x00')
			offset += ind_count
		
		for anim in k3result["animations"]:
			f.write(struct.pack("4H", anim["id"], len(anim["frames"]), anim["fps"], 0))
			
			for frame in anim["frames"]:
				for bone_idx in range(len(k3result["bones"])):
					f.write(struct.pack("4f", frame[bone_idx][0][0], frame[bone_idx][0][1], frame[bone_idx][0][2], 1))
					f.write(struct.pack("4f", *frame[bone_idx][1]))