import os
import sys

import bpy

script_dir = os.path.dirname(os.path.abspath(__file__))
utils_dir = os.path.join(script_dir, "../../blender_utils")
sys.path.append(utils_dir)

from utils import bake_model, clean_unused, export_ig_object, import_obj_folder

#############################################
# Parse command line arguments
#############################################


def get_arg(argv, flag, default=None):
    if flag in argv:
        return argv[argv.index(flag) + 1]
    return default


should_bake = "--bake" in sys.argv

axis = ["X", "Y", "Z", "-X", "-Y", "-Z"]
import_axis_up = get_arg(sys.argv, "--up", default="Z")
if import_axis_up not in axis:
    raise ValueError("Axis up not supported: {} (should be among X,Y,Z,-X,-Y,-Z)".format(import_axis_up))

import_axis_forward = get_arg(sys.argv, "--forward", default="X")
if import_axis_forward not in axis:
    raise ValueError("Axis forward not supported: {} (should be among X,Y,Z,-X,-Y,-Z)".format(import_axis_forward))

source_dir = get_arg(sys.argv, "--source_dir")
if source_dir is None:
    raise ValueError("Source directory not specified.")
dest_dir = get_arg(sys.argv, "--dest_dir")
if dest_dir is None:
    raise ValueError("Destination directory not specified.")
os.makedirs(dest_dir, exist_ok=True)

model_id = os.path.basename(source_dir)

#############################################
# Importing obj files from source dir
#############################################
for on in bpy.context.scene.objects.keys():
    obj = bpy.context.scene.objects[on]
    bpy.data.objects.remove(obj)
clean_unused()

import_obj_folder(model_id, source_dir, up=import_axis_up, forward=import_axis_forward)

#############################################
# Optional UV Unwrapping
# This only needed if baking will be performed
#############################################
if should_bake:
    uv_unwrapped = True
    for o in bpy.context.scene.objects:
        if not o.data.uv_layers:
            uv_unwrapped = False
    if not uv_unwrapped:
        bpy.ops.object.mode_set(mode="OBJECT")
        vl = bpy.context.view_layer
        bpy.ops.object.select_all(action="DESELECT")
        for on in bpy.context.scene.objects.keys():
            obj = bpy.context.scene.objects[on]
            new_uv = bpy.context.scene.objects[on].data.uv_layers.new(name="obj_uv")
            vl.objects.active = obj
            obj.select_set(True)
        bpy.ops.object.editmode_toggle()
        bpy.ops.mesh.select_all(action="SELECT")
        bpy.ops.uv.smart_project(angle_limit=66, island_margin=0.02)
        bpy.context.tool_settings.mesh_select_mode = (False, False, True)
        bpy.ops.object.mode_set(mode="OBJECT")


#############################################
# Export models
#############################################
export_ig_object(dest_dir, save_material=not should_bake)


#############################################
# Optional Texture Baking
#############################################
if should_bake:
    mat_dir = os.path.join(dest_dir, "material")
    os.makedirs(mat_dir, exist_ok=True)

    # bpy.ops.wm.open_mainfile(filepath=blend_path)
    # import_ig_object(model_root, import_mat=True)
    for obj in bpy.context.scene.objects:
        obj.select_set(True)
        bpy.context.view_layer.objects.active = obj
    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.join()

    channels = {
        "DIFFUSE": (2048, 32),
        "ROUGHNESS": (1024, 16),
        "METALLIC": (1024, 16),
        "NORMAL": (1024, 16),
    }
    bake_model(mat_dir, channels, overwrite=True)

bpy.ops.wm.quit_blender()

[(7, 'os.path.join', 'os.path.join', (['script_dir', '"""../../blender_utils"""'], {}), False, 'import os\n'), (8, 'sys.path.append', 'sys.path.append', (['utils_dir'], {}), False, 'import sys\n'), (40, 'os.makedirs', 'os.makedirs', (['dest_dir'], {'exist_ok': '(True)'}), False, 'import os\n'), (42, 'os.path.basename', 'os.path.basename', (['source_dir'], {}), False, 'import os\n'), (47, 'bpy.context.scene.objects.keys', 'bpy.context.scene.objects.keys', ([], {}), False, 'import bpy\n'), (50, 'utils.clean_unused', 'clean_unused', ([], {}), False, 'from utils import bake_model, clean_unused, export_ig_object, import_obj_folder\n'), (52, 'utils.import_obj_folder', 'import_obj_folder', (['model_id', 'source_dir'], {'up': 'import_axis_up', 'forward': 'import_axis_forward'}), False, 'from utils import bake_model, clean_unused, export_ig_object, import_obj_folder\n'), (82, 'utils.export_ig_object', 'export_ig_object', (['dest_dir'], {'save_material': '(not should_bake)'}), False, 'from utils import bake_model, clean_unused, export_ig_object, import_obj_folder\n'), (108, 'bpy.ops.wm.quit_blender', 'bpy.ops.wm.quit_blender', ([], {}), False, 'import bpy\n'), (6, 'os.path.abspath', 'os.path.abspath', (['__file__'], {}), False, 'import os\n'), (49, 'bpy.data.objects.remove', 'bpy.data.objects.remove', (['obj'], {}), False, 'import bpy\n'), (89, 'os.path.join', 'os.path.join', (['dest_dir', '"""material"""'], {}), False, 'import os\n'), (90, 'os.makedirs', 'os.makedirs', (['mat_dir'], {'exist_ok': '(True)'}), False, 'import os\n'), (97, 'bpy.ops.object.select_all', 'bpy.ops.object.select_all', ([], {'action': '"""SELECT"""'}), False, 'import bpy\n'), (98, 'bpy.ops.object.join', 'bpy.ops.object.join', ([], {}), False, 'import bpy\n'), (106, 'utils.bake_model', 'bake_model', (['mat_dir', 'channels'], {'overwrite': '(True)'}), False, 'from utils import bake_model, clean_unused, export_ig_object, import_obj_folder\n'), (64, 'bpy.ops.object.mode_set', 'bpy.ops.object.mode_set', ([], {'mode': '"""OBJECT"""'}), False, 'import bpy\n'), (66, 'bpy.ops.object.select_all', 'bpy.ops.object.select_all', ([], {'action': '"""DESELECT"""'}), False, 'import bpy\n'), (67, 'bpy.context.scene.objects.keys', 'bpy.context.scene.objects.keys', ([], {}), False, 'import bpy\n'), (72, 'bpy.ops.object.editmode_toggle', 'bpy.ops.object.editmode_toggle', ([], {}), False, 'import bpy\n'), (73, 'bpy.ops.mesh.select_all', 'bpy.ops.mesh.select_all', ([], {'action': '"""SELECT"""'}), False, 'import bpy\n'), (74, 'bpy.ops.uv.smart_project', 'bpy.ops.uv.smart_project', ([], {'angle_limit': '(66)', 'island_margin': '(0.02)'}), False, 'import bpy\n'), (76, 'bpy.ops.object.mode_set', 'bpy.ops.object.mode_set', ([], {'mode': '"""OBJECT"""'}), False, 'import bpy\n')]