import bpy
import math
import sys
import struct

bpy.ops.object.mode_set(mode = 'OBJECT')

obj = None
mesh = None
flat = False
if 'mesh' in bpy.data.objects:
	obj = bpy.data.objects['mesh']
	print('Using smooth normals')
if 'mesh_flat' in bpy.data.objects:
	obj = bpy.data.objects['mesh_flat']
	print('Using flat normals')
	flat = True
mesh = obj.data

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

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

armature = None
bone_names = None
if 'armature' in bpy.data.objects:
	armature = bpy.data.objects['armature']
	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:
	vertex_format |= COLOR_MASK
if texcoords:
	vertex_format |= TEXCOORD_MASK
if armature:
	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()

if flat:
	for p in mesh.loop_triangles:
		for li in p.loops:
			i = len(vertex_coords)
			v = mesh.vertices[mesh.loops[li].vertex_index]
			vertex_coords.append((v.co.x, v.co.y, v.co.z))
			vertex_normals.append((p.normal.x, p.normal.y, p.normal.z))
			if colors:
				vertex_colors.append(tuple(colors[li].color))
			if texcoords:
				vertex_texcoords.append(tuple(texcoords[li].uv))
			if armature:
				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)
				weights = list(sorted(weights, key=lambda p: -p[1]))[:2]
				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[vi] = weights
			indices.append(i)

else:
	for v in mesh.vertices:
		vertex_coords.append((v.co.x, v.co.y, v.co.z))
		vertex_normals.append((v.normal.x, v.normal.y, v.normal.z))

	if colors:
		vertex_colors = [None] * len(vertex_coords)
	if texcoords:
		vertex_texcoords = [None] * len(vertex_coords)
	if armature:
		vertex_weights = [None] * len(vertex_coords)

	for p in mesh.loop_triangles:
		for li in p.loops:
			vi = mesh.loops[li].vertex_index
			indices.append(vi)
			v = mesh.vertices[vi]
			if colors:
				vertex_colors[vi] = tuple(colors[li].color)
			if texcoords:
				vertex_texcoords[vi] = tuple(texcoords[li].uv)
			if armature:
				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)
				weights = list(sorted(weights, key=lambda p: -p[1]))[:2]
				while len(weights) < 2:
					weights.append((0,0.0))
				W = sum(p[1] for p in weights)
				weights = [(i,w/W) for i,w in weights]
				vertex_weights[vi] = weights


assert (len(indices) % 3) == 0
assert len(vertex_coords) == len(vertex_normals)
if colors:
	assert len(vertex_coords) == len(vertex_colors)
if texcoords:
	assert len(vertex_coords) == len(vertex_texcoords)
if armature:
	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:
	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:
	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:
		attribs.append(vertex_colors[i])
	if texcoords:
		attribs.append(vertex_texcoords[i])
	if armature:
		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:
	for index, pose in enumerate(bpy.data.objects['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))

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

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