# Chris Xiong 2018
# Expat (MIT) License
bl_info={
	"name":"3D XML (3.0) import",
	"description":"Import 3D XML 3.0",
	"author":"Chris Xiong",
	"version":(0,1),
	"blender":(2,82,0),
	"category":"Import-Export",
	"support":"TESTING"
}
################################################################
# This addon enables blender to import 3D XML 3.0 documents.
# It partially implemented the 3D XML specification 3.0 from
# Dassault Systèmes. Testing was done with files exported
# from 3DVIA Virtools 5.0.
#
# The development took place solely on Linux. It may experience
# problems on other platforms.
#
# This addon started off as a fork of this script:
# https://www.blender.org/forum/viewtopic.php?t=18299
# Later it was rewritten from scratch to better reflect the
# specification.
#
# To report bugs, mail me a minimal 3D XML file that can
# reproduce the issue together with the details of the bug.
#
# References:
# 1) 3D XML Reference Documentation 3.0:
# http://media.3ds.com/3dxml/3dxmldoc/Reference_Guide/3DXML_Reference_Guide.html
# 2) 3D XML 3.0 XSD schema:
# http://media.3ds.com/3dxml/3dxmldoc/3DXML.xsd
# 3) 3D XML User's Guide 3.0:
# http://media.3ds.com/3dxml/3dxmldoc/3DXML_User_Guide.pdf
################################################################
import bpy,bmesh,bpy_extras,mathutils
import xml.etree.ElementTree as etree
import pathlib,zipfile,time,os,tempfile,math

NS="{http://www.3ds.com/xsd/3DXML}"

meshes=dict()
meshmat=dict()
materials=dict()
textures=dict()
texturefiles=dict()

unitfactor=0.01292
texdir="/Ballance/Textures/"

def load_textures(tree):
	txd=pathlib.Path(texdir)
	for tex in tree.findall(f".//{NS}Image"):
		txname=tex.attrib["name"]
		txp=[]
		if txd.is_dir():
			txp=[i for i in txd.iterdir() if i.stem.lower()==txname.lower()]
		if len(txp)<1 and tex.attrib["href"] is not None and tex.attrib["href"] in texturefiles:
			txp=[texturefiles[tex.attrib["href"]]]
		tx=None
		try:
			tx=bpy.data.images.load(str(txp[0]),check_existing=True)
		except IndexError:
			print(txname)
		textures[tex.attrib["id"]]=tx

def load_materials(tree):
	for mat in tree.findall(f".//{NS}GraphicMaterial"):
		mname=mat.attrib["name"]
		m=bpy.data.materials.new(mname)
		m.use_nodes=True
		for node in m.node_tree.nodes:
			m.node_tree.nodes.remove(node)
		bnode=m.node_tree.nodes.new(type="ShaderNodeBsdfPrincipled")
		inode=m.node_tree.nodes.new(type="ShaderNodeTexImage")
		mnode=m.node_tree.nodes.new(type="ShaderNodeOutputMaterial")
		try:
			inode.image=textures[mat.attrib["texture"].split(":")[-1]]
		except KeyError:
			pass
		m.node_tree.links.new(inode.outputs[0],bnode.inputs[0])
		m.node_tree.links.new(bnode.outputs[0],mnode.inputs[0])
		m.diffuse_color=[float(mat.find(f"{NS}Diffuse").attrib[i])for i in ["red","green","blue"]]+[float(mat.attrib["diffuseCoef"])]
		m.specular_color=[float(mat.find(f"{NS}Specular").attrib[i])for i in ["red","green","blue"]]
		m.specular_intensity=float(mat.attrib["specularCoef"])
		materials[mat.attrib["id"]]=m

def load_meshes(tree):
	for rep in tree.findall(f".//{NS}Representation"):
		rid=rep.attrib["id"]
		verts=unflatten(rep.find(f".//{NS}Positions").text,float,3,unitfactor)
		normals=unflatten(rep.find(f".//{NS}Normals").text,float,3)
		uvs=unflatten(rep.find(f".//{NS}TextureCoordinates").text,float,2)
		faces=[]
		facemat=[]
		matslots=[]
		defmat=-1
		facesel=rep.find(f".//{NS}Faces")
		if facesel.find(f"{NS}SurfaceAttributes/{NS}MaterialApplication/{NS}MaterialId") is not None:
			defmat=0
			matslots.append(facesel.find(f"./{NS}SurfaceAttributes/{NS}MaterialApplication/{NS}MaterialId").text)
		for face in facesel.findall(f"{NS}Face"):
			fmat=defmat
			if face.find(f".//{NS}MaterialId") is not None:
				matp=-1
				try:
					matp=matslots.index(face.find(f".//{NS}MaterialId").text)
				except ValueError:
					matp=len(matslots)
					matslots.append(face.find(f".//{NS}MaterialId").text)
				fmat=matp
			if "triangles" in face.attrib:
				faces.extend(unflatten(face.attrib["triangles"],int,3))
				facemat.extend([fmat]*len(unflatten(face.attrib["triangles"],int,3)))
			if "fans" in face.attrib:
				faces.extend(unflatten(face.attrib["fans"],int,4))
				facemat.extend([fmat]*len(unflatten(face.attrib["fans"],int,4)))
		meshmat[rid]=matslots
		create_mesh(verts,faces,facemat,normals,uvs,rid)

def load_objects(tree):
	rr=tree.findall(f".//{NS}ReferenceRep[@format='TESSELLATED']")
	sr=set()
	for ref in rr:
		n=ref.attrib["name"]
		sr.add(n[:n.rfind('_')])
	ti3d=tree.findall(f".//{NS}Instance3D")
	i3d=[i3 for i3 in ti3d if i3.attrib["name"][:i3.attrib["name"].rfind('_')] in sr]
	for ref,i3 in zip(rr,i3d):
		meshid=ref.attrib["associatedFile"].split(":")[-1]
		objname=ref.attrib["name"]
		objname=objname[:objname.rfind('_')]
		mat=list(map(float,i3.find(f"./{NS}RelativeMatrix").text.split(' ')))
		obj=bpy.data.objects.new(objname,meshes[meshid])
		_wmat=mathutils.Matrix()
		for r in range(0,3):
			_wmat[r]=[mat[i] for i in range(r,12,3)]
		for m in meshmat[meshid]:
			obj.data.materials.append(materials[m])
		obj.matrix_world=_wmat
		scn=bpy.context.scene
		scn.collection.objects.link(obj)
		obj.select_set(True)

	bpy.ops.object.shade_smooth()

def create_mesh(verts,faces,facemat,norms,uvs,meshidx):
	if len(uvs)<len(verts):
		uvs.append([uvs[0]]*(len(verts)-len(uvs)))
	meshname=f"Mesh_{meshidx}"
	mesh=bmesh.new()
	for vert,norm in zip(verts,norms):
		v=mesh.verts.new(vert)
		v.normal=norm
	mesh.verts.ensure_lookup_table()
	mesh.verts.index_update()
	for i,m in zip(faces,facemat):
		f=[mesh.verts[j]for j in i]
		try:
			nf=mesh.faces.new(f)
			if m!=-1:nf.material_index=m
		except ValueError:
			pass
	uv=mesh.loops.layers.uv.new()
	for face in mesh.faces:
		for lp in face.loops:
			lp[uv].uv=mathutils.Vector(uvs[lp.vert.index])
	msh=bpy.data.meshes.new(meshname)
	if usesmooth:
		msh.use_auto_smooth=True
		msh.auto_smooth_angle=smoothangle
	mesh.to_mesh(msh)
	meshes[meshidx]=msh
	mesh.free()

def unflatten(seq,func,step,fac=1):
	seq=seq.replace(' ',',')
	seq=seq.split(',')
	seq=list(map(func,seq))
	nseq=len(seq)
	return [tuple(seq[i+j]*fac for j in range(0,step))for i in range(0,nseq,step)]

def read(filename):
	time1=time.time()
	if zipfile.is_zipfile(filename):
		zf=zipfile.ZipFile(filename,"r")
		try:
			member=zf.namelist().index("Root.3dxml")
		except ValueError:
			raise RuntimeError("not a 3D XML 3.0 file!")
		filename=zf.open("Root.3dxml")
		for im in zf.namelist():
			if im.endswith(".tga"):
				tf=tempfile.NamedTemporaryFile(suffix=".tga",delete=False)
				tf.write(zf.open(im).read())
				tf.flush()
				texturefiles[im]=tf.name
				tf.close()
	else:
		print("Warning: the file will be treated as a bare XML document.")
	tree=etree.parse(filename)
	load_textures(tree)
	load_materials(tree)
	load_meshes(tree)
	load_objects(tree)
	bpy.ops.file.pack_all()
	for fn,tfn in texturefiles.items():
		os.remove(tfn)
	texturefiles.clear()
	time2=time.time()
	print("Total import time is: %.2f seconds."%(time2-time1))

class ImportDialog(bpy.types.Operator,bpy_extras.io_utils.ImportHelper):
	bl_idname="object.tdxml_import"
	bl_label="3DXML Import"
	filter_glob=bpy.props.StringProperty(default='*.3dxml',options={'HIDDEN'})

	pt=bpy.props.StringProperty(name="Texture path",default=texdir)
	uf=bpy.props.FloatProperty(name="Unit factor",default=unitfactor,min=0)
	us=bpy.props.BoolProperty(name="Use auto smooth instead of the normal data in the 3DXML file",default=False)
	aa=bpy.props.FloatProperty(name="Auto smooth angle",min=0,max=180,default=90)

	def execute(self,context):
		global texdir,unitfactor,usesmooth,smoothangle
		texdir=self.pt
		unitfactor=self.uf
		usesmooth=self.us
		smoothangle=self.aa/180*math.pi
		read(self.filepath)
		return {'FINISHED'}

class _3DXMLImport(bpy.types.Operator):
	"""Import a 3DXML 3.0 file"""
	bl_idname="import_mesh.tdxml"
	bl_label="Import 3DXML 3.0..."
	bl_options={'UNDO'}
	filename_ext=".3dxml"

	def execute(self,context):
		bpy.ops.object.tdxml_import('INVOKE_DEFAULT')
		return {'FINISHED'}

def menu_func_import_tdxml(self,context):
	self.layout.operator(_3DXMLImport.bl_idname,text="3DXML 3.0 (.3dxml)")

def register():
	bpy.utils.register_class(ImportDialog)
	bpy.utils.register_class(_3DXMLImport)
	bpy.types.TOPBAR_MT_file_import.append(menu_func_import_tdxml)

def unregister():
	bpy.utils.unregister_class(_3DXMLImport)
	bpy.utils.unrigister_class(ImportDialog)

if __name__=="__main__":
	register()