# Blender addon to export an object as a lib.rs/chair mesh file
# To use, select object(s) and export with settings.
# Texture name in the file will be either:
# 1. custom "chair_texture" property on material
# 2. the material's name
# Based on the Raw mesh exporter distributed with blender

bl_info = {
	"name": "Chair Mesh Exporter",
	"author": "deltanedas",
	"blender": (3, 0, 0),
	"location": "File > Export > Chair Mesh (.chr)",
	# Crate verion 0.2.3
	"version": (0, 2, 7),
	"category": "Import-Export"
}

# For normal lookup
import math
# For writing floats
import struct

def face_to_tris(cfg, face, mesh, matrix, vertices, vertex_ids, colours, colour_ids, indices):
	if cfg.export_uvs:
		uvs = mesh.uv_layers.active.data

	if cfg.export_colours:
		cols = mesh.vertex_colors.active.data

	for n in range(3):
		verti = face.vertices[n]
		loopi = face.loops[n]
		vert = mesh.vertices[verti]
		loop = mesh.loops[loopi]

		if cfg.export_colours:
			col = (cols[loopi].color)[:]
			# really inefficient but this is pyton so
			colkey = str(col)
			if colkey in colour_ids:
				coli = colour_ids[colkey]
			else:
				# new colour, push to list
				coli = len(colours)
				colours.append(col)
				colour_ids[colkey] = coli
		else:
			coli = None

		if cfg.export_normals:
			normals = (loop.normal)[:]
		else:
			normals = None

		if cfg.export_uvs:
			uv = (uvs[loopi].uv)[:]
		else:
			uv = None

		v = [(matrix @ vert.co)[:], coli, normals, uv]
		# really inefficient but this is pyton so
		key = str(v)
		if key in vertex_ids:
			i = vertex_ids[key]
		else:
			# new vertex, push to list
			i = len(vertices)
			vertices.append(v)
			vertex_ids[key] = i

		indices.append(i)

def process_scene(cfg, context):
	colours = []
	colour_ids = {}
	textures = []
	texture_ids = {}
	objects = []
	indices = []

	scene = context.scene
	for obj in context.selected_objects:
		me = None
		if cfg.apply_modifiers or obj.type != "MESH":
			try:
				me = obj.to_mesh()
			except:
				me = None
			is_tmp_mesh = True
		else:
			me = obj.data
			is_tmp_mesh = False

		if me is None:
			continue

		if not me.loop_triangles and me.polygons:
			me.calc_loop_triangles()

		if cfg.export_normals:
			me.calc_normals_split()

		# TODO: vertex groups split object into its own thing
		if cfg.export_uvs:
			# get texture name from custom property or use the material name
			mat = obj.active_material
			if "chair_texutre" in mat:
				texture = mat["chair_texture"]
			else:
				texture = mat.name

			if texture in texture_ids:
				texturei = texture_ids[texture]
			else:
				texturei = len(textures)
				textures.append(texture)
				texture_ids[texture] = texturei
		else:
			texturei = None

		vertices = []
		vertex_ids = {}
		matrix = obj.matrix_world.copy()
		for face in me.loop_triangles:
			face_to_tris(cfg, face, me, matrix, vertices, vertex_ids, colours, colour_ids, indices)

		min_pos = [
			min(vertices, key=lambda v: v[0][0])[0][0],
			min(vertices, key=lambda v: v[0][1])[0][1],
			min(vertices, key=lambda v: v[0][2])[0][2]
		]
		max_pos = [
			max(vertices, key=lambda v: v[0][0])[0][0],
			max(vertices, key=lambda v: v[0][1])[0][1],
			max(vertices, key=lambda v: v[0][2])[0][2]
		]
		scale = [
			max_pos[0] - min_pos[0],
			max_pos[1] - min_pos[1],
			max_pos[2] - min_pos[2]
		]
		objects.append((vertices, scale, min_pos, texturei))

		# Remove meshes that are created solely for exporting
		if is_tmp_mesh:
			obj.to_mesh_clear()

	return colours, textures, objects, indices

def index_size(vertices):
	if vertices < 256:
		return 1
	if vertices < 65536:
		return 2
	return 4

def write_vec2_float(f, v):
	f.write(bytes(struct.pack("<f", v[0])))
	f.write(bytes(struct.pack("<f", 1.0 - v[1])))

def write_vec2_byte(f, v):
	x = v[0]
	y = v[1]
	while x < 0.0:
		x += 1.0
	while y < 0.0:
		y += 1.0
	while x > 1.0:
		x -= 1.0
	while y > 1.0:
		y -= 1.0
	x = int(x * 128)
	# flip y axis
	y = 128 - int(y * 128)
	f.write(bytes([x, y]))

def write_vec3_float(f, v):
	# <f = little-endian float
	f.write(bytes(struct.pack("<f", v[0])))
	# Y = up in chair/wgpu, Z = up in blender so swap them
	f.write(bytes(struct.pack("<f", v[2])))
	f.write(bytes(struct.pack("<f", v[1])))

def write_vec3_byte(f, v, scale, offset):
	# use range 0-128 to allow for even divisions
	x = 128 * (v[0] - offset[0]) / scale[0]
	# Y = up in chair/wgpu, Z = up in blender so swap them
	y = 128 * (v[2] - offset[2]) / scale[2]
	z = 128 * (v[1] - offset[1]) / scale[1]
	f.write(bytes([int(x), int(y), int(z)]))

tau = math.pi * 2
def index_normal(n):
	if n == 0:
		# exact mapping of 0 to 0.0 with no error
		return 0

	# asin fairly slow but this is pyton so doesnt matter
	return 128 - int(math.asin(n) * 255 / tau)

def write_vec3_sin(f, v):
	x = index_normal(v[0])
	# Y = up in chair/wgpu, Z = up in blender so swap them
	y = index_normal(v[2])
	z = index_normal(v[1])
	f.write(bytes([x, y, z]))

def write_model(cfg, colours, textures, objects, indices):
	f = open(cfg.filepath, "wb")
	# Chair file header
	f.write(b"CHAIR")
	# Format version
	f.write(bytes([0]))

	# MeshFeatures
	features = 0
	if cfg.export_colours:
		features |= 0b0001
	if cfg.export_normals:
		features |= 0b0010
	if cfg.export_uvs:
		features |= 0b0100
	if cfg.export_skeleton:
		features |= 0b1000
	if cfg.short_pos:
		features |= 0b00010000
	if cfg.short_normals:
		features |= 0b00100000
	if cfg.short_uvs:
		features |= 0b01000000
	if cfg.short_weights:
		features |= 0b10000000
	f.write(bytes([features]))

	if cfg.export_colours:
		f.write((len(colours)).to_bytes(2, "little"))
		for col in colours:
			for chan in col:
				f.write(bytes([int(chan * 255)]))
	if len(colours) > 255:
		colour_size = 2
	else:
		colour_size = 1

	if cfg.export_uvs:
		if len(textures) > 255:
			raise "Too many textures"

		f.write(bytes([len(textures)]))
		for texture in textures:
			name = texture.encode("utf-8")
			count = len(name)
			if count > 255:
				raise "Texture name " + texture + " is too long"

			f.write(bytes([count]))
			f.write(name)

	# TODO: export skeleton

	f.write(bytes([len(objects)]))
	for obj in objects:
		# Vertices count for this object
		f.write(len(obj[0]).to_bytes(4, "little"))

		scale = obj[1]
		offset = obj[2]
		if cfg.short_pos:
			write_vec3_float(f, scale)
			write_vec3_float(f, offset)
			write_pos = lambda f, v: write_vec3_byte(f, v, scale, offset)
		else:
			write_pos = write_vec3_float

		if cfg.short_normals:
			write_normals = write_vec3_sin
		else:
			write_normals = write_vec3_float

		if cfg.export_uvs:
			texture = obj[3]
			f.write(bytes([texture]))
		if cfg.short_uvs:
			write_uvs = write_vec2_byte
		else:
			write_uvs = write_vec2_float

		# Write data
		for v in obj[0]:
			write_pos(f, v[0])
			if cfg.export_colours:
				f.write(v[1].to_bytes(colour_size, "little"))
			if cfg.export_normals:
				write_normals(f, v[2])
			if cfg.export_uvs:
				write_uvs(f, v[3])
			# TODO: bone weights

	# Indices
	f.write(len(indices).to_bytes(4, "little"))
	indsize = index_size(sum([len(obj[0]) for obj in objects]))
	if cfg.flip_winding:
		for i in range(0, int(len(indices) / 3)):
			n = i * 3
			f.write(indices[n + 1].to_bytes(indsize, "little"))
			f.write(indices[n].to_bytes(indsize, "little"))
			f.write(indices[n + 2].to_bytes(indsize, "little"))
	else:
		for i in indices:
			f.write(i.to_bytes(indsize, "little"))

	f.close()
	return

import bpy
from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty, BoolProperty
from bpy.types import Operator

class ModelExporter(Operator, ExportHelper):
	"""Export geometry to chair mesh format (lib.rs/chair)"""
	bl_idname = "export.chair"
	bl_label = "Export Chair Mesh"

	filename_ext = ".chr"
	filter_glob: StringProperty(
		default = "*.chr",
		options = {"HIDDEN"}
	)

	apply_modifiers: BoolProperty(
		name = "Apply Modifiers",
		description = "Use transformed mesh data from each object",
		default = True
	)

	export_colours: BoolProperty(
		name = "Export Colours",
		description = "Save vertex colours in the exported mesh",
		default = False
	)

	export_normals: BoolProperty(
		name = "Export Normals",
		description = "Save vertex normals in the exported mesh",
		default = True
	)

	export_uvs: BoolProperty(
		name = "Export UVs",
		description = "Save vertex texture coordinates (and texture names) in the exported mesh",
		default = True
	)

	export_skeleton: BoolProperty(
		name = "Export Skeleton",
		description = "Save skeleton and vertex bones/weights in the exported mesh",
		default = False
	)

	short_pos: BoolProperty(
		name = "Short Positions",
		description = "Use bytes -64 to +64 instead of floats to lossily encode vertex positions.\nAlso stores scale + offset of objects",
		default = False
	)

	short_normals: BoolProperty(
		name = "Short Normals",
		description = "Use a byte sin lookup table instead of floats to lossily encode vertex normals",
		default = False
	)

	short_uvs: BoolProperty(
		name = "Short UVs",
		description = "Use bytes 0 to 128 instead of floats to lossily encode vertex texture coordinates",
		default = False
	)

	short_weights: BoolProperty(
		name = "Short Weights",
		description = "Use bytes -64 to +64 instead of floats to lossily encode vertex weights",
		default = False
	)

	flip_winding: BoolProperty(
		name = "Flip Winding",
		description = "Flip vertex winding from CCW to CW",
		default = False
	)

	def execute(self, context):
		colours, textures, objects, indices = process_scene(self, context)
		write_model(self, colours, textures, objects, indices)
		return {"FINISHED"}

def menu_export(self, context):
	self.layout.operator(ModelExporter.bl_idname, text = "Chair Mesh (.chr)")

def register():
	bpy.utils.register_class(ModelExporter)
	bpy.types.TOPBAR_MT_file_export.append(menu_export)

def unregister():
	bpy.utils.unregister_class(ModelExporter)
	bpy.types.TOPBAR_MT_file_export.remove(menu_export)

if __name__ == "__main__":
	register()