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, 12, 7, 59), 'os.path.join', 'os.path.join', ({(7, 25, 7, 35): 'script_dir', (7, 37, 7, 58): '"""../../blender_utils"""'}, {}), "(script_dir, '../../blender_utils')", False, 'import os\n'), ((8, 0, 8, 26), 'sys.path.append', 'sys.path.append', ({(8, 16, 8, 25): 'utils_dir'}, {}), '(utils_dir)', False, 'import sys\n'), ((40, 0, 40, 36), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((42, 11, 42, 39), 'os.path.basename', 'os.path.basename', ({(42, 28, 42, 38): 'source_dir'}, {}), '(source_dir)', False, 'import os\n'), ((47, 10, 47, 42), 'bpy.context.scene.objects.keys', 'bpy.context.scene.objects.keys', ({}, {}), '()', False, 'import bpy\n'), ((50, 0, 50, 14), 'utils.clean_unused', 'clean_unused', ({}, {}), '()', False, 'from utils import bake_model, clean_unused, export_ig_object, import_obj_folder\n'), ((52, 0, 52, 87), 'utils.import_obj_folder', 'import_obj_folder', (), '', False, 'from utils import bake_model, clean_unused, export_ig_object, import_obj_folder\n'), ((82, 0, 82, 57), 'utils.export_ig_object', 'export_ig_object', (), '', False, 'from utils import bake_model, clean_unused, export_ig_object, import_obj_folder\n'), ((108, 0, 108, 25), 'bpy.ops.wm.quit_blender', 'bpy.ops.wm.quit_blender', ({}, {}), '()', False, 'import bpy\n'), ((6, 29, 6, 54), 'os.path.abspath', 'os.path.abspath', ({(6, 45, 6, 53): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((49, 4, 49, 32), 'bpy.data.objects.remove', 'bpy.data.objects.remove', ({(49, 28, 49, 31): 'obj'}, {}), '(obj)', False, 'import bpy\n'), ((89, 14, 89, 48), 'os.path.join', 'os.path.join', ({(89, 27, 89, 35): 'dest_dir', (89, 37, 89, 47): '"""material"""'}, {}), "(dest_dir, 'material')", False, 'import os\n'), ((90, 4, 90, 39), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((97, 4, 97, 46), 'bpy.ops.object.select_all', 'bpy.ops.object.select_all', (), '', False, 'import bpy\n'), ((98, 4, 98, 25), 'bpy.ops.object.join', 'bpy.ops.object.join', ({}, {}), '()', False, 'import bpy\n'), ((106, 4, 106, 49), 'utils.bake_model', 'bake_model', (), '', False, 'from utils import bake_model, clean_unused, export_ig_object, import_obj_folder\n'), ((64, 8, 64, 46), 'bpy.ops.object.mode_set', 'bpy.ops.object.mode_set', (), '', False, 'import bpy\n'), ((66, 8, 66, 52), 'bpy.ops.object.select_all', 'bpy.ops.object.select_all', (), '', False, 'import bpy\n'), ((67, 18, 67, 50), 'bpy.context.scene.objects.keys', 'bpy.context.scene.objects.keys', ({}, {}), '()', False, 'import bpy\n'), ((72, 8, 72, 40), 'bpy.ops.object.editmode_toggle', 'bpy.ops.object.editmode_toggle', ({}, {}), '()', False, 'import bpy\n'), ((73, 8, 73, 48), 'bpy.ops.mesh.select_all', 'bpy.ops.mesh.select_all', (), '', False, 'import bpy\n'), ((74, 8, 74, 68), 'bpy.ops.uv.smart_project', 'bpy.ops.uv.smart_project', (), '', False, 'import bpy\n'), ((76, 8, 76, 46), 'bpy.ops.object.mode_set', 'bpy.ops.object.mode_set', (), '', False, 'import bpy\n')]