import bpy
import math
import sys
import struct

model_name = sys.argv[sys.argv.index('--') + 1]
filename = sys.argv[sys.argv.index('--') + 2]

obj = None
mesh = None
if model_name in bpy.data.objects:
	obj = bpy.data.objects[model_name]
else:
	raise RuntimeError("Model " + model_name + " not found in file")

mesh = obj.data

colors = None
if len(mesh.vertex_colors) > 0:
	colors = mesh.vertex_colors[0].data
	print('Found vertex colors')

texcoords = None
if len(mesh.uv_layers) > 0:
	texcoords = mesh.uv_layers[0].data
	print('Found texture coordinates:', texcoords)

armature = None
bone_names = None
if obj.parent is not None and type(obj.parent.data) is bpy.types.Armature:
	armature = obj.parent
	print('Found armature with {} bones'.format(len(armature.data.bones)))
	bone_names = [b.name for b in armature.data.bones]

POSITION_MASK = 1
NORMAL_MASK = 2
COLOR_MASK = 4
TEXCOORD_MASK = 8
WEIGHTS_MASK = 16

vertex_format = POSITION_MASK | NORMAL_MASK
if colors is not None:
	vertex_format |= COLOR_MASK
if texcoords is not None:
	vertex_format |= TEXCOORD_MASK
if armature is not None:
	vertex_format |= WEIGHTS_MASK
print("Using vertex format", format(vertex_format, '05b'))

vertex_coords = []
vertex_normals = []
vertex_colors = []
vertex_texcoords = []
vertex_weights = []

indices = []

mesh.calc_loop_triangles()

smooth_index = {}

def append_vertex(v, n):
	vertex_coords.append((v.co.x, v.co.y, v.co.z))
	vertex_normals.append((n.x, n.y, n.z))
	if colors is not None:
		vertex_colors.append(tuple(colors[li].color))
	if texcoords is not None:
		vertex_texcoords.append((texcoords[li].uv[0], 1.0 - texcoords[li].uv[1]))
	if armature is not None:
		weights = []
		for g in v.groups:
			name = obj.vertex_groups[g.group].name
			if name in bone_names:
				weights.append((bone_names.index(name), g.weight))
		for g, w in weights:
			if g >= len(armature.data.bones):
				print("Invalid vertex group index {} for armature weights".format(g), file=sys.stderr)
				sys.exit(1)
		while len(weights) < 2:
			weights.append((0,0.0))
		weights = list(sorted(weights, key=lambda p: -p[1]))[:2]
		W = sum(p[1] for p in weights)
		weights = [(i,w/W) for i,w in weights]
		vertex_weights.append(weights)

for p in mesh.loop_triangles:
	for li in p.loops:
		vi = mesh.loops[li].vertex_index
		v = mesh.vertices[vi]
		n = p.normal
		if p.use_smooth:
			i = None
			if vi not in smooth_index:
				i = len(vertex_coords)
				smooth_index[vi] = i
				append_vertex(v, n)
			else:
				i = smooth_index[vi]

			indices.append(i)
		else:
			i = len(vertex_coords)
			append_vertex(v, n)
			indices.append(i)


assert (len(indices) % 3) == 0
assert len(vertex_coords) == len(vertex_normals)
if colors is not None:
	assert len(vertex_coords) == len(vertex_colors)
if texcoords is not None:
	assert len(vertex_coords) == len(vertex_texcoords)
if armature is not None:
	assert len(vertex_coords) == len(vertex_weights)


class Uint8:
	def __init__(self, value):
		self.value = int(value)

	def __str__(self):
		return str(self.value)

	def __repr__(self):
		return str(self.value)


if colors is not None:
	for i in range(len(vertex_colors)):
		c = 0
		for k in (3, 2, 1, 0):
			c = (c << 8) | int(max(0, min(255, pow(vertex_colors[i][k], 2.2) * 255)))
		vertex_colors[i] = c

if armature is not None:
	for i in range(len(vertex_weights)):
		gs = []
		for g, w in vertex_weights[i]:
			gs.append(Uint8(g))

		ws = []
		for g, w in vertex_weights[i]:
			ws.append(Uint8(w * 255.0))
		ws[0].value += 255 - sum(w.value for w in ws)

		assert len(gs) == 2
		assert len(ws) == 2
		vertex_weights[i] = tuple(gs + ws)

vertices = []
for i in range(len(vertex_coords)):
	attribs = [vertex_coords[i], vertex_normals[i]]
	if colors is not None:
		attribs.append(vertex_colors[i])
	if texcoords is not None:
		attribs.append(vertex_texcoords[i])
	if armature is not None:
		attribs.append(vertex_weights[i])
	vertices.append(tuple(attribs))

print(len(vertices), 'vertices')
print(len(indices), 'indices')

if armature is not None:
	bones = []
	for b in armature.data.bones:
		p = None
		if b.parent is None:
			p = -1
		else:
			p = list(armature.data.bones).index(b.parent)
		bones.append((p, tuple(b.head_local), tuple(map(tuple, b.matrix_local.to_3x3()))))

poses = []
if armature is not None and armature.pose_library is not None:
	for index, pose in enumerate(armature.pose_library.pose_markers):
		print("Found pose", pose.name)

		# [rotation, scale, translation]
		data = []
		for i in range(len(bone_names)):
			data.append([[1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [0.0, 0.0, 0.0]])

		marker = armature.pose_library.pose_markers[index]
		frame = marker.frame
		action =  bpy.data.actions[armature.pose_library.name]
		for g in action.groups:
			if not g.name in bone_names:
				continue
			bone_index = bone_names.index(g.name)
			for ch in g.channels:
				value = ch.evaluate(float(frame))
				if ch.data_path.find('rotation_quaternion') != -1:
					data[bone_index][0][ch.array_index] = value
				elif ch.data_path.find('scale') != -1:
					data[bone_index][1][ch.array_index] = value
				elif ch.data_path.find('location') != -1:
					data[bone_index][2][ch.array_index] = value

		new_data = []
		for d in data:
			rotation = tuple([d[0][1], d[0][2], d[0][3], d[0][0]])
			scale = d[1][0]
			translation = tuple(d[2])
			new_data.append((rotation, scale, translation))
		poses.append((pose.name, new_data))

def to_bytes(obj):
	if type(obj) == Uint8:
		return struct.pack('<B', obj.value)
	if type(obj) == int:
		if obj < 0:
			return struct.pack('<i', obj)
		else:
			return struct.pack('<I', obj)
	if type(obj) == float:
		return struct.pack('<f', obj)
	if type(obj) == tuple:
		res = bytes()
		for x in obj:
			res += to_bytes(x)
		return res
	if type(obj) == list:
		res = to_bytes(len(obj))
		for x in obj:
			res += to_bytes(x)
		return res
	if type(obj) == str:
		b = obj.encode('utf-8')
		return struct.pack('<I', len(b)) + b
	print("Type", type(obj), "not supported")
	sys.exit(1)

SECTION_MESH  = 1
SECTION_BONES = 2
SECTION_POSE  = 3

data = to_bytes((SECTION_MESH, vertex_format, vertices, indices))

if armature:
	data += to_bytes((SECTION_BONES, bones))

for name, pdata in poses:
	data += to_bytes((SECTION_POSE, name, pdata))

with open(filename, 'wb') as f:
	f.write(data)