ap-canvas-creation-module-old/03_blender/sd_blender/__init__.py

# ----------------------------------------------------------
# File __init__.py
# ----------------------------------------------------------

# Imports
# Check if this add-on is being reloaded
if "bpy" in locals():
    # reloading .py files
    import bpy
    import importlib

    # from . import zs_renderscene as zsrs
    # importlib.reload(zsrs)


# or if this is the first load of this add-on
else:
    import bpy
    import json
    from pathlib import Path
    from dataclasses import dataclass
    from mathutils import Euler, Vector, Quaternion, Matrix
    import math
    import os
    import base64
    import numpy as np
    import shutil

    # from . import zs_renderscene as zsrs  # noqa


# Addon info
bl_info = {
    "name": "ZS Stable Diffusion Connection V0.0.1",
    "author": "Sergiu <sergiu@zixelise.com>",
    "Version": (0, 0, 1),
    "blender": (4, 00, 0),
    "category": "Scene",
    "description": "Stable Diffusion Connection",
}


# -------------------------------------------------------------------
# Functions
# -------------------------------------------------------------------


@dataclass
class AssetData:
    path: str
    collection_name: str
    type: str


product_asset = AssetData(
    path="sample_scene/01_Products", collection_name="01_Products", type="product"
)
element_asset = AssetData(
    path="sample_scene/02_Elements", collection_name="02_Elements", type="asset"
)
basic_shapes_asset = AssetData(
    path="sample_scene/03_BasicShapes",
    collection_name="03_BasicShapes",
    type="shape",
)


def convert_base64_string_to_object(base64_string):
    bytes = base64.b64decode(base64_string)
    string = bytes.decode("ascii")

    return json.loads(string)
    return string


def load_scene():

    print("Loading Scene")
    # load scene data

    scene_data = load_scene_data()
    print(scene_data)
    # create parent collections
    create_parent_collections(product_asset.collection_name)
    create_parent_collections(element_asset.collection_name)
    create_parent_collections(basic_shapes_asset.collection_name)
    create_parent_collections("05_Cameras")
    # append products
    products_data = load_objects_data(scene_data, product_asset.type)
    for index, product in enumerate(products_data):
        append_objects(product, index, product_asset)
    # append elements
    elements_data = load_objects_data(scene_data, element_asset.type)
    for index, product in enumerate(elements_data):
        append_objects(product, index, element_asset)
    # append shapes
    basic_shapes_data = load_objects_data(scene_data, basic_shapes_asset.type)
    for index, product in enumerate(basic_shapes_data):
        append_objects(product, index, basic_shapes_asset)

    # set lighting
    set_environment(scene_data)

    # set camera
    create_cameras(scene_data)

    # rename outputs

    set_output_paths(
        "D://Git//ap-canvas-creation-module//03_blender//sd_blender//sample_scene//Renders",
        scene_data["project_id"],
    )

    # setup compositing
    set_cryptomatte_objects("01_Products", "mask_product")


def invert_id_name(json_data):
    for obj in json_data["scene"]["objects"]:
        obj["name"], obj["id"] = obj["id"], obj["name"]
    return json_data


def load_scene_data():
    # print("Loading Scene Data")
    # # load scene data

    # # to be replaced with actual data
    # # open scene_info.json
    # script_path = Path(__file__).resolve()
    # scene_data_path = script_path.parent / "sample_scene" / "scene_info.json"
    # with scene_data_path.open() as file:
    #     scene_data = json.load(file)
    #     print(scene_data)
    # return scene_data

    # load scene data
    print("Loading Scene Data")
    if bpy.context.scene.load_local_DB:
        loaded_scene_data = bpy.context.scene.config_string
        # check if loaded_scene_data is base64 encoded
        if loaded_scene_data.startswith("ey"):
            scene_data = convert_base64_string_to_object(loaded_scene_data)
        else:
            scene_data = json.loads(loaded_scene_data)

    else:
        scene_data = json.loads(bpy.context.scene.shot_info_ai)

    # invert_scene_data = invert_id_name(scene_data)

    return scene_data


def load_objects_data(scene_data, object_type: str):
    print("Loading Assets Data")
    # load assets data
    objects_data = []
    for object in scene_data["scene"]["objects"]:
        if object["group_type"] == object_type:
            # get additional object data by id and combine with object data
            object_data = get_object_data_by_id(object["id"])
            if object_data:
                # temporary fix
                # object_data = get_object_data_by_id(object["name"])
                object.update(object_data)
                objects_data.append(object)
            else:
                print("Object not found in database", object["id"], object["name"])
    return objects_data

    # to be replaced with actual data


def get_object_data_by_id(object_id: str):
    print("Getting Object Data")
    # open assets_database.json
    script_path = Path(__file__).resolve()
    assets_data_path = script_path.parent / "sample_scene" / "assets_database.json"

    with assets_data_path.open() as file:
        assets_data = json.load(file)
    # get object data by id
    for object in assets_data:
        if object["id"] == object_id:
            return object


def get_hdri_data_by_id(object_id: str):
    print("Getting HDRI Data")
    # open assets_database.json
    script_path = Path(__file__).resolve()
    assets_data_path = script_path.parent / "sample_scene" / "lighting_database.json"

    with assets_data_path.open() as file:
        assets_data = json.load(file)
    # get object data by id
    for object in assets_data:
        if object["id"] == object_id:
            return object


def append_objects(asset_info, index, asset_data: AssetData):
    print("Appending Objects")

    blendFileName = asset_info["name"]

    # visibleLayersJSONName = productInfo["Version"]

    # activeSwitchMaterials = json.dumps(productInfo["ActiveMaterials"])

    collectionName = blendFileName + "_" + str(index)

    append_active_layers(collectionName, asset_info, asset_data)

    # replace_switch_materials(shotInfo, productInfo["ActiveMaterials"])

    link_collection_to_collection(
        asset_data.collection_name, bpy.data.collections[collectionName]
    )


def append_active_layers(newCollectionName, product_info, asset_data: AssetData):
    utility_collections = [
        "MaterialLibrary",
        "Animation_Controller",
        "Animation_Rig",
        "Animation_Target",
    ]

    # visible_layers = utility_collections + product_info["VisibleLayers"]
    visible_layers = utility_collections

    script_path = Path(__file__).resolve().parent
    filePath = str(
        script_path
        / asset_data.path
        / product_info["name"]
        / "BLEND"
        / (product_info["name"] + ".blend")
    )

    print(filePath)

    # delete all objects and collections from product collection

    create_parent_collections(newCollectionName)

    # append active collections
    with bpy.data.libraries.load(filePath) as (data_from, data_to):
        data_to.collections = []
        for name in data_from.collections:

            if name in visible_layers:
                data_to.collections.append(name)
            if "NonConfigurable" in name:
                data_to.collections.append(name)

    # link appended colections to newCollection
    for collection in data_to.collections:
        # try:
        #     bpy.context.scene.collection.children.link(collection)
        # except:
        #     print(collection)
        link_collection_to_collection(newCollectionName, collection)

        # hide utility collections
        for utilityCollectionName in utility_collections:
            if utilityCollectionName in collection.name:
                # rename utility collections
                collection.name = newCollectionName + "_" + utilityCollectionName
                hide_collection(collection)

            # need to redo this in the future
            if "Animation_Target" in collection.name:

                # set the x location
                collection.objects[0].location.x = product_info["properties"][
                    "transform"
                ]["position"][0]
                # set the y location
                collection.objects[0].location.y = -product_info["properties"][
                    "transform"
                ]["position"][2]
                # set the z location
                collection.objects[0].location.z = product_info["properties"][
                    "transform"
                ]["position"][1]

                # collection.objects[0].location = product_info["properties"][
                #     "transform"
                # ]["position"]

                # collection.objects[0].rotation_euler = product_info["properties"][
                #     "transform"
                # ]["rotation"]

                # set object rotation

                # rotation_in_degrees = product_info["properties"]["transform"][
                #     "rotation"
                # ]

                # rotation_in_degrees[0] = rotation_in_degrees[0] + 90

                # # set object rotation in euler from radians
                # collection.objects[0].rotation_euler = (
                #     math.radians(rotation_in_degrees[0]),
                #     math.radians(rotation_in_degrees[2]),
                #     math.radians(rotation_in_degrees[1]),
                # )

                rotation_data_euler = product_info["properties"]["transform"][
                    "rotationEuler"
                ]

                collection.objects[0].rotation_euler = (
                    math.radians(
                        rotation_data_euler[0]
                    ),  # Add 90 degrees to the first element
                    math.radians(rotation_data_euler[2]),
                    math.radians(rotation_data_euler[1]),
                )

                # # set rotation using quaternion

                # rotation_quat_data = product_info["properties"]["transform"]["rotation"]

                # converted_quat = Quaternion(
                #     [
                #         rotation_quat_data[3],
                #         rotation_quat_data[0],
                #         rotation_quat_data[1],
                #         rotation_quat_data[2],
                #     ]
                # )

                # # Define the camera correction quaternion (from GLTF file)
                # rotation_correction = Quaternion((2**0.5 / 2, 2**0.5 / 2, 0.0, 0.0))

                # # Apply the camera correction to the converted quaternion
                # corrected_quat = rotation_correction @ converted_quat

                # # Apply the corrected quaternion to the camera's rotation
                # collection.objects[0].rotation_mode = "QUATERNION"
                # collection.objects[0].rotation_quaternion = rotation_correction

                # set object scale
                collection.objects[0].scale = product_info["properties"]["transform"][
                    "scale"
                ]

    # make all objects in collection local
    for obj in bpy.data.collections[newCollectionName].all_objects:
        if obj.type == "MESH":
            obj.make_local()
            obj.data.make_local()

    # remove duplicated material slots
    mats = bpy.data.materials

    for obj in bpy.data.collections[newCollectionName].all_objects:
        for slot in obj.material_slots:
            part = slot.name.rpartition(".")
            if part[2].isnumeric() and part[0] in mats:
                slot.material = mats.get(part[0])


def create_parent_collections(group_name: str):

    if collection_exists(group_name):

        remove_collection_and_objects(group_name)
    else:
        create_collection(group_name)


def set_environment(scene_data):
    # Find the group named NG_Canvas_Background in the Blender world
    world = bpy.context.scene.world

    if world is not None:
        # Get the node group named NG_Canvas_Background
        node_group = world.node_tree.nodes.get("NG_Canvas_Background")

        if node_group is not None:
            # Set the group's properties from scene_data
            node_group.inputs[0].default_value = float(
                scene_data["scene"]["environment"]["lighting"]["rotation"]
            )
            node_group.inputs[1].default_value = float(
                scene_data["scene"]["environment"]["lighting"]["intensity"]
            )

            if scene_data["scene"]["environment"]["lighting"]["visible"] == True:
                node_group.inputs[2].default_value = 1.0
            else:
                node_group.inputs[2].default_value = 0.0

            node_group.inputs[3].default_value = (
                scene_data["scene"]["environment"]["background"]["color"]["r"],
                scene_data["scene"]["environment"]["background"]["color"]["g"],
                scene_data["scene"]["environment"]["background"]["color"]["b"],
                1.0,
            )

            hdri_data = get_hdri_data_by_id(
                scene_data["scene"]["environment"]["lighting"]["id"]
            )

            # go inside the node group, and find the image texture node, and set the image to the one from the scene data
            for node in node_group.node_tree.nodes:
                if node.name == "environment_map":
                    node.image = bpy.data.images.load(
                        str(
                            Path(__file__).resolve().parent
                            / "sample_scene"
                            / "05_Lighting"
                            / "HDRI"
                            / (hdri_data["name"] + ".exr")
                        )
                    )
        else:
            print("Group 'NG_Canvas_Background' not found")


def calculate_focal_length(fov, film_height):
    # Convert FOV from degrees to radians
    fov_rad = math.radians(fov)

    # Calculate the focal length
    focal_length = (film_height / 2) / math.tan(fov_rad / 2)

    return focal_length


def quaternion_multiply(q1, q2):
    """
    Multiplies two quaternions.
    q1 and q2 are arrays or lists of length 4.
    Returns the resulting quaternion as a NumPy array.
    """
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
    return np.array([w, x, y, z])


def create_cameras(scene_data):
    # # Get the 05_Cameras collection, or create it if it doesn't exist
    collection_name = "05_Cameras"
    if collection_name in bpy.data.collections:
        collection = bpy.data.collections[collection_name]
    else:
        return

    # Assuming `scene_data` and `collection` are already defined
    for camera_data in scene_data["scene"]["cameras"]:
        # Create a new camera object
        bpy.ops.object.camera_add()

        # Get the newly created camera
        camera = bpy.context.object

        # Set the camera's name
        camera.name = camera_data["name"]

        # Set the camera's position
        position = camera_data["properties"]["transform"]["position"]
        camera.location.x = position[0]
        camera.location.y = -position[2]
        camera.location.z = position[1]

        # Set the camera's rotation
        # # euler
        # rotation_euler = camera_data["properties"]["transform"]["rotation"]
        # rotation_euler = Euler(
        #     (
        #         math.radians(rotation_euler[0] + 90),
        #         math.radians(-rotation_euler[2]),
        #         math.radians(rotation_euler[1]),
        #     ),
        #     "XYZ",
        # )
        # quaternion
        rotation_quat_data = camera_data["properties"]["transform"]["rotation"]
        # rotation_quat = (
        #     rotation_quat[3],
        #     rotation_quat[0],
        #     rotation_quat[1],
        #     rotation_quat[2],
        # )

        # new_quat = Quaternion((2**0.5 / 2, 2**0.5 / 2, 0.0, 0.0))

        # current_quat = Quaternion(rotation_quat)

        # result_quat = current_quat @ new_quat

        # rotation_quat = (result_quat.w, result_quat.x, result_quat.y, result_quat.z)

        rotation_quat = [
            rotation_quat_data[3],
            rotation_quat_data[0],
            rotation_quat_data[1],
            rotation_quat_data[2],
        ]  # Example quaternion

        # new_quat = [2**0.5 / 2, 2**0.5 / 2, 0.0, 0.0]

        # result_quat = quaternion_multiply(rotation_quat, new_quat)

        # Example quaternion from GLTF: [x, y, z, w]
        gltf_quat = [0.0, 0.707, 0.0, 0.707]

        # Convert the GLTF quaternion to Blender space (Y-up to Z-up)
        converted_quat = Quaternion(
            [
                rotation_quat_data[3],
                rotation_quat_data[0],
                rotation_quat_data[1],
                rotation_quat_data[2],
            ]
        )

        # Define the camera correction quaternion (from GLTF file)
        camera_correction = Quaternion((2**0.5 / 2, 2**0.5 / 2, 0.0, 0.0))

        # Apply the camera correction to the converted quaternion
        corrected_quat = camera_correction @ converted_quat

        # Apply the corrected quaternion to the camera's rotation
        camera.rotation_mode = "QUATERNION"
        camera.rotation_quaternion = corrected_quat

        # camera.rotation_mode = "QUATERNION"
        # camera.rotation_quaternion = rotation_quat

        # Apply the local rotation to the camera

        # Set the camera's lens properties
        lens = camera_data["properties"]["lens"]
        type_mapping = {
            "PERSPECTIVE": "PERSP",
            "ORTHOGRAPHIC": "ORTHO",
            "PANORAMIC": "PANO",
        }

        camera.data.lens = lens["focalLength"]

        camera.data.type = type_mapping.get(lens["type"].upper(), "PERSP")

        camera.data.clip_start = lens["near"]
        camera.data.clip_end = lens["far"]

        # Add the camera to the 05_Cameras collection
        collection.objects.link(camera)

        # Unlink the camera from the scene collection
        # check all objects in scene collection, if camera is there, unlink it
        for obj in bpy.context.scene.collection.objects:
            if obj.name == camera.name:
                bpy.context.scene.collection.objects.unlink(camera)

        # Set the camera as the active camera if "active" is true
        if camera_data["properties"]["active"]:
            bpy.context.scene.camera = camera


def set_output_paths(base_path, project_name):

    # check if folder exist, if not create it
    folder_path = base_path + "//" + project_name
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    # Get the current scene
    scene = bpy.context.scene

    # Check if the scene has a compositor node tree
    if scene.node_tree:
        # Iterate over all nodes in the node tree
        for node in scene.node_tree.nodes:
            # Check if the node is an output node
            if node.type == "OUTPUT_FILE":
                # Set the base path of the output node
                node.base_path = folder_path
                # Iterate over all file slots of the output node
                # for file_slot in node.file_slots:
                #     # Set the path of the file slot
                #     file_slot.path = project_name


def set_cryptomatte_objects(collection_name, node_name):
    # Get all objects in the specified collection
    objects = bpy.data.collections[collection_name].all_objects

    # Convert the objects to a list
    object_names = [obj.name for obj in objects]

    print(object_names)

    # Get the current scene
    scene = bpy.context.scene

    # Check if the scene has a compositor node tree
    if scene.node_tree:
        # Iterate over all nodes in the node tree
        for node in scene.node_tree.nodes:
            # Check if the node is a Cryptomatte node with the specified name
            if node.type == "CRYPTOMATTE_V2" and node.name == node_name:
                # Set the Matte ID objects of the node
                node.matte_id = ",".join(object_names)


# -------------------------------------------------------------------
# Utilities
# -------------------------------------------------------------------
def get_all_objects_in_collection(collection):
    """Recursively get all objects in the given collection and its subcollections."""
    objects = []
    for obj in collection.objects:
        objects.append(obj)
    for subcollection in collection.children:
        objects.extend(get_all_objects_in_collection(subcollection))
    return objects


def remove_collection_and_objects(collection_name):
    oldObjects = list(bpy.data.collections[collection_name].all_objects)
    for obj in oldObjects:
        bpy.data.objects.remove(obj, do_unlink=True)

    old_collection = bpy.data.collections[collection_name]
    if old_collection is not None:
        old_collection_names = get_subcollection_names(old_collection)
    else:
        print("Collection not found.")

    # print line break
    print("-----------------------------------------------------------------")
    print(old_collection_names)
    print("-----------------------------------------------------------------")
    for old_collection_name in old_collection_names:
        for collection in bpy.data.collections:
            if collection.name == old_collection_name:
                bpy.data.collections.remove(collection)

    bpy.ops.outliner.orphans_purge(
        do_local_ids=True, do_linked_ids=True, do_recursive=True
    )


def get_subcollection_names(collection):
    subcollection_names = []

    for child in collection.children:
        subcollection_names.append(child.name)
        subcollection_names.extend(get_subcollection_names(child))

    return subcollection_names


def select_objects_in_collection(collection):
    """Recursively select all objects in the given collection and its subcollections."""
    for obj in collection.objects:
        obj.select_set(True)
    for subcollection in collection.children:
        select_objects_in_collection(subcollection)


def link_collection_to_collection(parentCollectionName, childCollection):

    if bpy.context.scene.collection.children:
        parentCollection = bpy.context.scene.collection.children.get(
            parentCollectionName
        )
        # Add it to the main collection

        # childCollection = bpy.context.scene.collection.children.get(childCollection)

        parentCollection.children.link(childCollection)
        # if child collection is in scene collection unlink it
        if bpy.context.scene.collection.children.get(childCollection.name):
            bpy.context.scene.collection.children.unlink(childCollection)
        # bpy.context.scene.collection.children.unlink(childCollection)


# link collection to collection
def link_collection_to_collection_old(parentCollectionName, childCollection):

    if bpy.context.scene.collection.children:
        parentCollection = bpy.context.scene.collection.children.get(
            parentCollectionName
        )
        # Add it to the main collection

        try:
            childCollection = bpy.data.collections[childCollection]
        except:
            print("Collection not found.")
            return

        parentCollection.children.link(childCollection)
        bpy.context.scene.collection.children.unlink(childCollection)


# function that checks if a collection exists
def collection_exists(collection_name):
    return collection_name in bpy.data.collections


# function that creates a new collection and adds it to the scene
def create_collection(collection_name):
    new_collection = bpy.data.collections.new(collection_name)
    bpy.context.scene.collection.children.link(new_collection)


def hide_collection(collection):
    collection.hide_render = True
    collection.hide_viewport = True


def check_if_selected_objects_have_parent(self):
    for obj in bpy.context.selected_objects:
        if obj.parent is None:
            message = f"Object {obj.name} has no parent"
            self.report({"ERROR"}, message)
            return False
    return True


def add_rig_controller_to_selection(self):
    # add "Rig_Controller_Main" to the selection
    has_controller_object = False
    for obj in bpy.data.objects:
        if obj.name == "Rig_Controller_Main":
            if obj.hide_viewport:
                self.report({"ERROR"}, "Rig_Controller_Main is hidden")
                return has_controller_object
            obj.select_set(True)
            # if object is not visible, make it visible

            has_controller_object = True
            return has_controller_object

    if not has_controller_object:
        message = f"Rig_Controller_Main not found"
        self.report({"ERROR"}, message)
        print("Rig_Controller_Main not found")

    return has_controller_object


def select_objects_in_collection(collection):
    """Recursively select all objects in the given collection and its subcollections."""
    for obj in collection.objects:
        obj.select_set(True)
    for subcollection in collection.children:
        select_objects_in_collection(subcollection)


def check_if_object_has_principled_material(obj):
    for slot in obj.material_slots:

        # if more than one slot is principled, return
        for node in slot.material.node_tree.nodes:
            # print(node.type)
            if node.type == "BSDF_PRINCIPLED":

                return True
            else:
                print("Object has no principled material", obj.name)
                return False

    return True


def unhide_all_objects():
    # show all objects using operator

    bpy.ops.object.hide_view_clear()


# -------------------------------------------------------------------
# Scene optimization
# -------------------------------------------------------------------
def export_non_configurable_to_fbx():
    # Get the current .blend file path

    blend_filepath = bpy.data.filepath
    if not blend_filepath:
        print("Save the .blend file first.")
        return

    # Get the parent directory of the .blend file
    blend_dir = os.path.dirname(blend_filepath)
    parent_dir = os.path.dirname(blend_dir)
    blend_filename = os.path.splitext(os.path.basename(blend_filepath))[0]

    # Create the FBX export path
    fbx_export_path = os.path.join(parent_dir, "FBX", f"{blend_filename}.fbx")

    # Ensure the FBX directory exists
    os.makedirs(os.path.dirname(fbx_export_path), exist_ok=True)

    # Deselect all objects
    bpy.ops.object.select_all(action="DESELECT")

    # Select all objects in the NonConfigurable collection
    collection_name = "NonConfigurable"
    if collection_name in bpy.data.collections:
        collection = bpy.data.collections[collection_name]
        for obj in collection.objects:
            obj.select_set(True)
    else:
        print(f"Collection '{collection_name}' not found.")
        return


def export_scene_to_fbx(self):
    # Ensure the .blend file is saved
    if not bpy.data.is_saved:
        print("Save the .blend file first.")
        self.report({"ERROR"}, "Save the .blend file first.")
        return

    # Get the current .blend file path
    blend_filepath = bpy.data.filepath
    if not blend_filepath:
        print("Unable to get the .blend file path.")
        return

    # Get the parent directory of the .blend file
    blend_dir = os.path.dirname(blend_filepath)
    parent_dir = os.path.dirname(blend_dir)
    blend_filename = os.path.splitext(os.path.basename(blend_filepath))[0]

    # Create the FBX export path
    fbx_export_path = os.path.join(parent_dir, "FBX", f"{blend_filename}.fbx")

    # Ensure the FBX directory exists
    os.makedirs(os.path.dirname(fbx_export_path), exist_ok=True)

    # unhide all objects
    unhide_all_objects()
    # Deselect all objects
    bpy.ops.object.select_all(action="DESELECT")

    # Select all objects in the NonConfigurable collection and its subcollections
    collection_name = "NonConfigurable"
    if collection_name in bpy.data.collections:
        collection = bpy.data.collections[collection_name]
        select_objects_in_collection(collection)
    else:
        print(f"Collection '{collection_name}' not found.")
        return

        # check if all objects selected have a parent, if not return

    if not check_if_selected_objects_have_parent(self):
        return

    if not add_rig_controller_to_selection(self):
        return

    # Export selected objects to FBX
    bpy.ops.export_scene.fbx(
        filepath=fbx_export_path,
        use_selection=True,
        global_scale=1.0,
        apply_unit_scale=True,
        bake_space_transform=False,
        object_types={"MESH", "ARMATURE", "EMPTY"},
        use_mesh_modifiers=True,
        mesh_smooth_type="FACE",
        use_custom_props=True,
        # bake_anim=False,
    )

    print(f"Exported to {fbx_export_path}")


def export_scene_to_glb(self):
    # Ensure the .blend file is saved
    if not bpy.data.is_saved:
        print("Save the .blend file first.")
        self.report({"ERROR"}, "Save the .blend file first.")
        return

    # Get the current .blend file path
    blend_filepath = bpy.data.filepath
    if not blend_filepath:
        print("Unable to get the .blend file path.")
        return

    # Get the parent directory of the .blend file
    blend_dir = os.path.dirname(blend_filepath)
    parent_dir = os.path.dirname(blend_dir)
    blend_filename = os.path.splitext(os.path.basename(blend_filepath))[0]

    # Create the GLB export path
    glb_export_path = os.path.join(parent_dir, "WEB", f"{blend_filename}.glb")

    # Ensure the GLB directory exists
    os.makedirs(os.path.dirname(glb_export_path), exist_ok=True)

    # unhide all objects
    unhide_all_objects()
    # Deselect all objects
    bpy.ops.object.select_all(action="DESELECT")

    # Select all objects in the NonConfigurable collection and its subcollections
    collection_name = "Web"
    for collection in bpy.data.collections:
        if collection_name == collection.name:
            select_objects_in_collection(collection)
            if not check_if_selected_objects_have_parent(self):
                return
            # check if all objects selected have a parent, if not return
            if not add_rig_controller_to_selection(self):
                return

            # # for each selected objects, check if the the material is principled, if not return
            # for obj in bpy.context.selected_objects:
            #     if not check_if_object_has_principled_material(obj):
            #         message = f"Object {obj.name} has no principled material"
            #         self.report({"ERROR"}, message)
            #         return

            # Export selected objects to GLB
            bpy.ops.export_scene.gltf(
                filepath=glb_export_path,
                export_format="GLB",
                use_selection=True,
                export_apply=True,
                export_animations=False,
                export_yup=True,
                export_cameras=False,
                export_lights=False,
                export_materials="EXPORT",
                export_normals=True,
                export_tangents=True,
                export_morph=False,
                export_skins=False,
                export_draco_mesh_compression_enable=False,
                export_draco_mesh_compression_level=6,
                export_draco_position_quantization=14,
                export_draco_normal_quantization=10,
                export_draco_texcoord_quantization=12,
                export_draco_color_quantization=10,
                export_draco_generic_quantization=12,
                export_keep_originals=False,
                export_texture_dir="",
            )

            print(f"Exported to {glb_export_path}")

    else:
        print(f"Collection '{collection_name}' not found.")
        return


def copy_and_relink_textures(collection_name):
    # Ensure the target directory exists
    target_directory = bpy.path.abspath("//..")
    textures_directory = os.path.join(target_directory, "01_Textures")
    os.makedirs(textures_directory, exist_ok=True)

    # Get the collection
    collection = bpy.data.collections.get(collection_name)
    if not collection:
        print(f"Collection '{collection_name}' not found.")
        return

    # Iterate through each object in the collection
    objects = get_all_objects_in_collection(collection)
    for obj in objects:
        if obj.type != "MESH":
            continue

        # Iterate through each material of the object
        for mat_slot in obj.material_slots:
            material = mat_slot.material
            if not material:
                continue

            # Check if the material contains any image textures
            if material.use_nodes:
                for node in material.node_tree.nodes:
                    if node.type == "TEX_IMAGE":
                        image = node.image
                        if image:
                            # Copy the image texture to the target directory
                            src_path = bpy.path.abspath(image.filepath)
                            filename = os.path.basename(src_path)
                            dest_path = os.path.join(textures_directory, filename)

                            # Check if the source and destination paths are the same
                            if os.path.abspath(src_path) != os.path.abspath(dest_path):
                                # Check if the file already exists in the destination directory
                                if not os.path.exists(dest_path):
                                    shutil.copy(src_path, dest_path)
                                    print(
                                        f"File '{filename}' copied to '{textures_directory}'"
                                    )
                                else:
                                    print(
                                        f"File '{filename}' already exists in '{textures_directory}', relinking."
                                    )

                            # Relink the image texture to the new location
                            image.filepath = bpy.path.relpath(dest_path)
                            image.reload()

    print(f"Textures copied and relinked to '{textures_directory}'.")


def set_assets_render_output_paths():
    # Get the current scene
    scene = bpy.context.scene
    target_directory = bpy.path.abspath("//..")
    custom_folder = os.path.join(target_directory, "PNG")
    blender_file_name = os.path.splitext(bpy.path.basename(bpy.data.filepath))[0]

    components = blender_file_name.split("_")
    if len(components) != 7:
        raise ValueError(
            "Blender file name must be in the format 'Brand_AssetName_Year_ContentUsage_ContentType_AssetType_AssetNumber'"
        )

    brand = components[0]
    asset_name = components[1]
    year = components[2]
    content_usage = components[3]
    content_type = components[4]
    asset_type = components[5]
    asset_number = components[6]

    imag_asset_type = "IMG"

    # Construct the new .blend file name
    image_name = f"{brand}_{asset_name}_{year}_{content_usage}_{content_type}_{imag_asset_type}_{asset_number}_"

    # Ensure the scene has a compositor node tree
    if not scene.use_nodes:
        print("Scene does not use nodes.")
        return

    node_tree = scene.node_tree

    # Search for the ZS_Canvas_Output group node
    for node in node_tree.nodes:
        if node.type == "GROUP" and node.node_tree.name == "ZS_Canvas_Output":
            # Check inside the group for all file output nodes
            for sub_node in node.node_tree.nodes:
                if sub_node.type == "OUTPUT_FILE":
                    # Set the base_path to the custom folder
                    sub_node.base_path = custom_folder

                    # Set the path for each file output slot to the Blender file name

                    for output in sub_node.file_slots:
                        output.path = image_name

    print("Output paths set.")


# -------------------------------------------------------------------
# Operators
# -------------------------------------------------------------------
# load scene operator
class ZSSD_OT_LoadScene(bpy.types.Operator):
    bl_idname = "zs_sd_loader.load_scene"
    bl_label = "Load Scene"
    bl_description = "Load Scene"

    def execute(self, context):
        load_scene()
        return {"FINISHED"}


# canvas exporter operator
class ZSSD_OT_ExportAssets(bpy.types.Operator):
    bl_idname = "zs_canvas.export_assets"
    bl_label = "Export Assets"
    bl_description = "Export Scene Assets to FBX and GLB"

    def execute(self, context):
        copy_and_relink_textures("NonConfigurable")
        copy_and_relink_textures("Web")
        set_assets_render_output_paths()
        export_scene_to_fbx(self)
        export_scene_to_glb(self)
        return {"FINISHED"}


# parent class for panels
class ZSSDPanel:
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "ZS SD Loader"


# -------------------------------------------------------------------
#   Draw
# -------------------------------------------------------------------

# Panels


class ZSSD_PT_Main(ZSSDPanel, bpy.types.Panel):
    bl_label = "SD Loader"

    def draw(self, context):
        layout = self.layout

        scene = context.scene

        col = layout.column()

        self.is_connected = False

        col.label(text="Stable Diffusion Connection")

        col.prop(context.scene, "load_local_DB")

        col.prop(context.scene, "config_string")

        # load scene button

        col.operator("zs_sd_loader.load_scene", text="Load Scene")

        col.separator()

        # export assets button
        col.operator("zs_canvas.export_assets", text="Export Assets")


# modify after making products
blender_classes = [
    ZSSD_PT_Main,
    ZSSD_OT_LoadScene,
    ZSSD_OT_ExportAssets,
]


def register():

    # register classes
    for blender_class in blender_classes:
        bpy.utils.register_class(blender_class)

    bpy.types.Scene.shot_info_ai = bpy.props.StringProperty(
        name="Shot Info",
    )

    bpy.types.Scene.config_string = bpy.props.StringProperty(  # type: ignore
        name="Configuration String",
    )

    bpy.types.Scene.load_local_DB = bpy.props.BoolProperty(  # type: ignore
        name="Load Local DB",
    )

    # Has to be afqter class registering to correctly register property

    # register global properties

    # register list

    # list data

    # bpy.types.Scene.zs_product_list = bpy.props.CollectionProperty(
    #    type=ZS_Product_ListItem)

    # current item in list


def unregister():

    # unregister classes
    for blender_class in blender_classes:
        bpy.utils.unregister_class(blender_class)
    # unregister global properties
    del bpy.types.Scene.shot_info_ai
    del bpy.types.Scene.config_string
    del bpy.types.Scene.load_local_DB

    # unregister list items
    # del bpy.types.Scene.my_list
    # del bpy.types.Scene.product_product_list_index