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
    from pyquaternion import Quaternion

    # 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",
}


# -------------------------------------------------------------------
# Convert Data Functions
# -------------------------------------------------------------------
def convert_loc(x):
    return u * Vector([x[0], -x[2], x[1]])


def convert_quat(q):
    return Quaternion([q[3], q[0], -q[2], q[1]])


def convert_scale(s):
    return Vector([s[0], s[2], s[1]])


def local_rotation(obj, rotation_before, rot):
    """Appends a local rotation to vnode's world transform:
    (new world transform) = (old world transform) @ (rot)
    without changing the world transform of vnode's children.

    For correctness, rot must be a signed permutation of the axes
    (eg. (X Y Z)->(X -Z Y)) OR vnode's scale must always be uniform.
    """
    rotation_before = Quaternion((1, 0, 0, 0))
    obj.rotation_before @= rot

    # # Append the inverse rotation after children's TRS to cancel it out.
    # rot_inv = rot.conjugated()
    # for child in gltf.vnodes[vnode_id].children:
    #     gltf.vnodes[child].rotation_after = rot_inv @ gltf.vnodes[child].rotation_after


# -------------------------------------------------------------------
# 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"]

                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]),
                )

                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 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
        rotation = camera_data["properties"]["transform"]["rotation"]
        rotation_euler = Euler(
            (
                math.radians(rotation[0]),
                math.radians(rotation[2]),
                math.radians(rotation[1]),
            ),
            "XYZ",
        )

        # Apply the local rotation to the camera
        camera.rotation_euler = rotation_euler

        # Set the camera's lens properties
        lens = camera_data["properties"]["lens"]
        type_mapping = {
            "PERSPECTIVE": "PERSP",
            "ORTHOGRAPHIC": "ORTHO",
            "PANORAMIC": "PANO",
        }
        camera.data.type = type_mapping.get(lens["type"].upper(), "PERSP")
        camera.data.angle = math.radians(lens["fov"])
        camera.data.clip_start = lens["near"]
        camera.data.clip_end = lens["far"]

        # Add the camera to the 05_Cameras collection
        collection.objects.link(camera)
        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 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 = "WebGL"
    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

    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}")


# -------------------------------------------------------------------
# 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):
        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