#!/usr/bin/env python
# Obj2tmdl
# Converter for Wavefront OBJ => Treeki Model

import os
import struct
import sys
u32 = struct.Struct('>I')
u16 = struct.Struct('>H')
u8 = struct.Struct('>B')
f32 = struct.Struct('>f')

sys.path.append('/home/Treeki/Wii.py/Wii.py')
import Wii

class MaterialLib:
	class Material:
		pass
	
	
	def __init__(self, filename):
		self.current_mat = None
		self.materials = {}
		
		for line in open(filename, 'r'):
			self.parse_line(line)
	
	
	def parse_line(self, line):
		line = line.strip()
		if line == '' or line[0] == '#':
			return
		
		line = line.split()
		cmd = line[0]
		
		if cmd == 'newmtl':
			self.current_mat = MaterialLib.Material()
			self.materials[line[1]] = self.current_mat
			
		elif cmd == 'map_Ka':
			self.current_mat.texture = ' '.join(line[1:])
	
	
	def prepare_textures(self):
		# get a list of every texture file used
		self.texture_id = {}
		textures = []
		id = 0
		
		for mat in self.materials.values():
			if hasattr(mat, 'texture'):
				if mat.texture not in textures:
					textures.append(mat.texture)
					self.texture_id[mat.texture] = id
					id += 1
		
		
		if len(textures) == 0:
			self.tpl = None
			return
		
		# create a gxtexconv script file and an ID mapping
		#script = []
		#self.texture_id = {}
		#
		#for tex, id in zip(textures, range(len(textures))):
		#	script.append('<filepath=\"%s\" id=\"tex%d\" colfmt=5 />' % (tex.replace('\\','/'),id))
		#	self.texture_id[id] = tex
		#
		#print script
		#
		#open('texture_script_temp_0991.scf', 'w').write('\n'.join(script))
		#
		#os.system('gxtexconv -s texture_script_temp_0991.scf -o generated_texture_0991.tpl')
		
		# fix this later
		texture_files = [tex.replace('\\','/').replace(' ','\\ ').replace('.jpg','.png') for tex in textures]
		print 'zetsubou wpng RGB5A3 %s generated_texture_0991.tpl' % ' '.join(texture_files)
		os.system('zetsubou wpng RGB5A3 %s generated_texture_0991.tpl' % ' '.join(texture_files))
		self.tpl = open('generated_texture_0991.tpl', 'rb').read()
		
		#os.remove('texture_script_temp_0991.scf')
		os.remove('generated_texture_0991.tpl')



class ObjReader:
	class ObjShape:
		def __init__(self, name):
			self.faces = []
			self.shape_name = name
			self.material_name = None
	
	
	class ObjFace:
		__slots__ = ('vertices', 'texcoords', 'normals')
		
		def __init__(self, vertices, texcoords, normals):
			self.vertices = vertices
			self.texcoords = texcoords
			self.normals = normals
	
	
	def __init__(self):
		self.material_lib = None
		
		self.vertex_lists = []
		self.texcoord_lists = []
		self.normal_lists = []
		self.shapes = []
		
		self.current_shape = None
	
	
	def parse_line(self, line):
		line = line.strip()
		if line == '' or line[0] == '#':
			return
		
		line = line.split()
		cmd = line[0]
		
		if cmd == 'v':
			self.vertex_lists.append(map(float, line[1:]))
			
		elif cmd == 'vt':
			self.texcoord_lists.append(map(float, line[1:])[:2]) # chop off third coord
			
		elif cmd == 'vn':
			self.normal_lists.append(map(float, line[1:]))
			
		elif cmd == 'f':
			v, t, n = [], [], []
			
			for entry in line[1:]:
				entry_split = map(lambda x:-1 if x == '' else int(x), entry.split('/'))
				
				v.append(self.vertex_lists[entry_split[0] - 1])
				
				if entry_split[1] == -1:
					t.append([0.0, 0.0])
				else:
					t.append(self.texcoord_lists[entry_split[1] - 1])
				
				if entry_split[2] == -1:
					n.append([0.0, 0.0, 0.0])
				else:
					n.append(self.normal_lists[entry_split[2] - 1])
			
			self.current_shape.faces.append(ObjReader.ObjFace(v, t, n))
			
		#elif cmd == 'g':
		#	self.current_shape = ObjReader.ObjShape(line[1])
		#	self.shapes.append(self.current_shape)
			
		elif cmd == 'mtllib':
			self.material_lib = MaterialLib(' '.join(line[1:]))
			
		elif cmd == 'usemtl':
			self.current_shape = ObjReader.ObjShape(line[1])
			self.shapes.append(self.current_shape)
			self.current_shape.material_name = line[1]



# === Treeki Model Format ===
# Header:
#
# struct tmdl_header {
# 	u32 magic;				// always TMDL
# 	u32 version;			// currently 2
# 	u32 shape_count;
# };
# [[Followed by an array of u32s for each shape offset]]
#
# struct tmdl_shape {
# 	u32 dl_offset;			// must be aligned to 0x20
# 	u32 dl_size;			// must be aligned to 0x20
# 	u32 mat_offset;
# };
#
# struct tmdl_material {
# 	u32 texture_id;			// once the model is bound, points to id within tpl
# };
#
# dl_offset points to a GX display list which is used for
# rendering the model in question.


class TmdlWriter:
	def __init__(self):
		self.shapes = []
	
	
	def build_with_obj(self, obj):
		self.shapes = obj.shapes
		self.material_lib = obj.material_lib
		self.materials = obj.material_lib.materials
		
		self.material_lib.prepare_textures()
	
	
	def pack(self):
		offset = 0
		
		header = struct.pack('>4sI I', 'TMDL', 2, len(self.shapes))
		offset += len(header)
		
		# reserve space for the shape offsets, we'll add them later
		offset += (4 * len(self.shapes))
		
		# create the materials
		material_offs = {}
		material_data = ''
		
		for name,material in self.materials.iteritems():
			material_offs[name] = offset
			material_struct = self.build_material_struct(material)
			material_data += material_struct
			offset += len(material_struct)
		
		
		# create the shapes -- first, we'll build the DLs
		shape_display_lists = map(self.build_display_list_with_shape, self.shapes)
		
		# calculate offsets to every display list
		# shape struct is currently 12 bytes, change this if the size changes
		# also, make sure they're aligned to an offset of 0x20
		temp_offset = offset + (len(self.shapes) * 12)
		
		if temp_offset & 0x1f != 0:
			aligned_offset = (temp_offset + 0x20) & ~0x1f
			dl_start_padding = '\0' * (aligned_offset - temp_offset)
			temp_offset = aligned_offset
		else:
			dl_start_padding = ''
		
		shape_display_list_offsets = []
		for dl in shape_display_lists:
			shape_display_list_offsets.append(temp_offset)
			temp_offset += len(dl)
		
		
		# and now, create the shape structs themselves
		shape_offsets = []
		shape_data = ''
		
		for shape, dl, dl_offset in zip(self.shapes, shape_display_lists, shape_display_list_offsets):
			# first off, store the offset
			shape_offsets.append(offset)
			
			# now build the struct
			shape_struct = struct.pack('>III', dl_offset, len(dl), material_offs[shape.material_name])
			shape_data += shape_struct
			offset += len(shape_struct)
		
		
		# almost there!
		tmdl_bits = [header]
		tmdl_add = tmdl_bits.append
		
		for offs in shape_offsets:
			tmdl_add(u32.pack(offs))
		
		tmdl_add(material_data)
		tmdl_add(shape_data)
		tmdl_add(dl_start_padding)
		tmdl_bits += shape_display_lists
		
		return ''.join(tmdl_bits)
	
	
	def build_material_struct(self, material):
		texture_id = 0xFFFFFFFF
		if hasattr(material, 'texture'):
			texture_id = self.material_lib.texture_id[material.texture]
		
		return u32.pack(texture_id)
	
	
	def build_display_list_with_shape(self, shape):
		# http://hitmen.c02.at/files/yagcd/yagcd/chap5.html
		# assume vertex descriptor: [data is specified in this order]
		# GXClearVtxDesc()
		# GXSetVtxDesc(GX_VA_POS, GX_DIRECT)
		# GXSetVtxDesc(GX_VA_NRM, GX_DIRECT)
		# GXSetVtxDesc(GX_VA_TEX0, GX_DIRECT)
		
		vtx_num = reduce(lambda x,y:x+len(y.vertices), shape.faces, 0)
		
		dl_bits = []
		dl_add = dl_bits.append
		
		# commands: GX_QUADS = 0x80, GX_TRIANGLES = 0x90
		dl_add('\x90') # opcode
		dl_add(u16.pack(vtx_num))
		
		# now send over the data
		for face in shape.faces:
			for v, t, n in zip(face.vertices, face.texcoords, face.normals):
				for f in v:
					dl_add(f32.pack(f))
				
				for f in n:
					dl_add(f32.pack(f))
				
				for f in t:
					dl_add(f32.pack(f))
		
		# done!
		return self.build_display_list(dl_bits)
	
	
	def build_display_list(self, dl_bits):
		# assemble it
		dl = ''.join(dl_bits)
		
		# pad the display list to 0x20 bytes with GX_NOP (0x00)
		if len(dl) % 0x20 != 0:
			pad_count = ((len(dl) + 0x20) & ~0x1F) - len(dl)
			dl += '\x00' * pad_count
		
		return dl


#file = r'H:\ISOs\NSMBWii\testmush3'
#file = 'testmdl2'
#file = 'simple2'
file = sys.argv[1]
shortfn = file[file.rfind('/')+1:]

obj = ObjReader()
for line in open(file+'.obj'):
	obj.parse_line(line)

tmdl = TmdlWriter()
tmdl.build_with_obj(obj)

tmdl_file = tmdl.pack()
tpl_file = tmdl.material_lib.tpl

arc = Wii.U8()
arc['t3d'] = None
arc['t3d/mdl_%s.tmdl' % shortfn] = tmdl_file
if tpl_file: arc['t3d/tex_%s.tpl' % shortfn] = tpl_file
arc.dumpFile('%s.arc' % file)