# Blender script that imports an OBJ (3D model) exported from Alone in the Dark
# room viewer application (https://github.com/tigrouind/AITD-roomviewer), and
# fixes the materials so that it can be properly rendered/edited.

# NOTE: the noise texture is needed (get it from the Assets/Materials/Model/
# directory of the mentioned application). Put it in the same folder as this
# script so Blender can find it. It will also have to be present in the same
# folder as the output .blend file, unless you manually change the path or pack
# (embed) it into the .blend.

# Usage:
#     blender --background --python <script.py> -- -f <file.obj> [-o <file.blend>]


import argparse
import bpy
import os
import sys


USAGE_TEXT = "Usage:\n\tblender --background --python " + __file__ + " -- -f <file.obj> [-o <file.blend>]"
SUPPORTED_MATERIALS = {"shadeless", "noise", "glass", "metal_horizontal", "metal_vertical"}
NOISE_TEXTURE = "noise.png"


def main():
	argv = sys.argv

	# "--" (with spaces before and after) separates the Blender arguments from
	# the Python script ones. We recover these arguments (if there are any).
	if "--" not in argv:
		print(USAGE_TEXT)
		return
	else:
		argv = argv[argv.index("--") + 1:]

	# Use the handy ArgumentParser to do the job for us.
	argParser = argparse.ArgumentParser(description="Blender script that imports an OBJ (3D model) exported from Alone in the Dark room viewer application.")
	argParser.add_argument("-f", "--file", dest="inFile", type=str, required=True, help="The .obj file that will be read.")
	argParser.add_argument("-o", "--output", dest="outFile", type=str, required=False, help="The .blend file that will be saved (optional, default 'untitled.blend').", default="untitled.blend")
	args = argParser.parse_args(argv)

	# Part 1: clear the scene (remove all objects).
	bpy.ops.wm.read_homefile(use_empty=True)

	# Part 2: try to import the .obj file into Blender.
	# Optional interesting arguments for obj_import():
	#     global_scale (default 1.0)
	#     clamp_size   (default 0.0)
	#     forward_axis (default "NEGATIVE_Z")
	#     up_axis      (default "Y")
	# Check reference at https://docs.blender.org/api/current/bpy.ops.wm.html#bpy.ops.wm.obj_import
	result = str(bpy.ops.wm.obj_import(filepath=args.inFile, forward_axis="Z"))
	result = result.strip() # Remove leading and trailing spaces.
	if result != "{'FINISHED'}":
		print("There was an error trying to import the chosen OBJ. Check that it is a valid file.")
		return

	# Part 3: fix the materials.
	for material in bpy.data.materials:
		# We will only fix these materials we are interested in. We will avoid
		# messing with others (in case there are any).
		if material.name not in SUPPORTED_MATERIALS:
			continue
		material.use_nodes = True
		material.node_tree.nodes.clear()
		# SHADELESS (standard, plain color with no effects).
		if material.name == "shadeless":
			nodeVertexColor = material.node_tree.nodes.new("ShaderNodeVertexColor")
			nodeVertexColor.location = (0, 0)
			nodeOutput = material.node_tree.nodes.new("ShaderNodeOutputMaterial")
			nodeOutput.location = (200, 0)
			material.node_tree.links.new(nodeOutput.inputs["Surface"], nodeVertexColor.outputs["Color"])
		# NOISE (noisy pattern mixed with the plain color).
		elif material.name == "noise":
			nodeVertexColor = material.node_tree.nodes.new("ShaderNodeVertexColor")
			nodeVertexColor.location = (100, 150)
			nodeImageTexture = material.node_tree.nodes.new("ShaderNodeTexImage")
			try:
				nodeImageTexture.image = bpy.data.images.load(os.path.join(os.getcwd(), NOISE_TEXTURE))
			except Exception as e:
				print(f"Error loading '{NOISE_TEXTURE}': {e}")
				return
			nodeImageTexture.location = (0, 0)
			nodeMix = material.node_tree.nodes.new("ShaderNodeMix")
			nodeMix.data_type = "RGBA"
			nodeMix.blend_type = "MULTIPLY"
			nodeMix.clamp_factor = True
			nodeMix.inputs["Factor"].default_value = 0.5
			nodeMix.location = (350, 150)
			nodeOutput = material.node_tree.nodes.new("ShaderNodeOutputMaterial")
			nodeOutput.location = (550, 150)
			material.node_tree.links.new(nodeMix.inputs["A"], nodeVertexColor.outputs["Color"])
			material.node_tree.links.new(nodeMix.inputs["B"], nodeImageTexture.outputs["Color"])
			material.node_tree.links.new(nodeOutput.inputs["Surface"], nodeMix.outputs["Result"])
		# GLASS (surfaces with simple transparency and no refraction, where
		#     white -> fully transparent
		#     black -> fully opaque
		# and all other colors having an intermediate degree of transparency).
		elif material.name == "glass":
			nodeVertexColor = material.node_tree.nodes.new("ShaderNodeVertexColor")
			nodeVertexColor.location = (0, 0)
			nodeShader = material.node_tree.nodes.new("ShaderNodeBsdfTransparent")
			nodeShader.location = (200, 0)
			nodeOutput = material.node_tree.nodes.new("ShaderNodeOutputMaterial")
			nodeOutput.location = (400, 0)
			material.node_tree.links.new(nodeShader.inputs["Color"], nodeVertexColor.outputs["Color"])
			material.node_tree.links.new(nodeOutput.inputs["Surface"], nodeShader.outputs["BSDF"])
			material.blend_method = "BLEND"
		# METAL_HORIZONTAL / METAL_VERTICAL (metallic surfaces, with a gradient).
		elif material.name == "metal_horizontal" or material.name == "metal_vertical":
			nodeTexCoord = material.node_tree.nodes.new("ShaderNodeTexCoord")
			nodeTexCoord.location = (0, 0)
			nodeSeparateXYZ = material.node_tree.nodes.new("ShaderNodeSeparateXYZ")
			nodeSeparateXYZ.location = (200, 0)
			nodeMultiply = material.node_tree.nodes.new("ShaderNodeMath")
			nodeMultiply.operation = "MULTIPLY"
			nodeMultiply.inputs[1].default_value = 5.0
			nodeMultiply.location = (400, 0)
			nodePingPong = material.node_tree.nodes.new("ShaderNodeMath")
			nodePingPong.operation = "PINGPONG"
			nodePingPong.inputs[1].default_value = 1.0
			nodePingPong.location = (600, 0)
			nodeAdd = material.node_tree.nodes.new("ShaderNodeMath")
			nodeAdd.operation = "ADD"
			nodeAdd.inputs[1].default_value = 0.5
			nodeAdd.location = (800, 0)
			nodeVertexColor = material.node_tree.nodes.new("ShaderNodeVertexColor")
			nodeVertexColor.location = (800, 150)
			nodeGamma = material.node_tree.nodes.new("ShaderNodeGamma")
			nodeGamma.location = (1000, 100)
			nodeOutput = material.node_tree.nodes.new("ShaderNodeOutputMaterial")
			nodeOutput.location = (1200, 100)
			material.node_tree.links.new(nodeSeparateXYZ.inputs["Vector"], nodeTexCoord.outputs["Window"])
			if material.name == "metal_horizontal":
				material.node_tree.links.new(nodeMultiply.inputs[0], nodeSeparateXYZ.outputs["X"])
			elif material.name == "metal_vertical":
				material.node_tree.links.new(nodeMultiply.inputs[0], nodeSeparateXYZ.outputs["Y"])
			material.node_tree.links.new(nodePingPong.inputs[0], nodeMultiply.outputs["Value"])
			material.node_tree.links.new(nodeAdd.inputs[0], nodePingPong.outputs["Value"])
			material.node_tree.links.new(nodeGamma.inputs["Color"], nodeVertexColor.outputs["Color"])
			material.node_tree.links.new(nodeGamma.inputs["Gamma"], nodeAdd.outputs["Value"])
			material.node_tree.links.new(nodeOutput.inputs["Surface"], nodeGamma.outputs["Color"])

	# Part 4: set the 3D viewport shading mode to "material preview", to allow
	# a quick check upon opening the .blend file without having to do anything.
	if bpy.context.screen:
		for screenArea in bpy.context.screen.areas:
			if screenArea.type == "VIEW_3D":
				screenArea.spaces[0].shading.type = "MATERIAL"

	# Part 5: save the .blend result file.
	try:
		bpy.ops.wm.save_mainfile(filepath=args.outFile)
		# Blender does not allow to set relative paths before saving the file.
		# To prevent "noise texture not found" errors, we do it now and resave.
		bpy.ops.file.make_paths_relative()
		bpy.ops.wm.save_mainfile()
	except Exception as e:
		print(f"Error saving .blend file: {e}")


if __name__ == "__main__":
	main()