Untitled

import bpy
import uuid

from pathlib import Path
from ayon_blender.api import plugin, lib


SHADERBALL_PATH = Path("Z:/Dev/ayon-blender/client/ayon_blender/plugins/create/external_data/look_shaderball.blend")
SHADERBALL_OBJECT_NAME = "look_shaderBall"
PLACEHOLDER_MATERIAL_NAME = "mat_shaderBall_placeHolder"
AVALON_CONTAINERS = "AVALON_CONTAINERS"
AVALON_INSTANCES = "AVALON_INSTANCES"
DEBUG_VERSION = "1.1.1"

class CreateLook(plugin.BlenderCreator):
    """LookFile containing shading/materials."""

    identifier = "io.openpype.creators.blender.look"
    label = "Look"
    product_type = "look"
    icon = "paint-brush"

    create_as_asset_group = False  # LookFiles do not need a group like models

    def create(
        self, product_name: str, instance_data: dict, pre_create_data: dict
    ):
        """Create a LookFile asset instance in Blender."""
        instance = super().create(product_name, instance_data, pre_create_data)
        look_variant = instance_data.get("variant", "Main")  # Get Look variant from Ayon Creator
        print(f"[DEBUG v{DEBUG_VERSION}] Creating LookFile instance for variant: {look_variant}...")
        self.create_shaderballs_for_look(look_variant)
        self.imprint_ayon_metadata(instance, look_variant)
        return instance

    def imprint_ayon_metadata(self, instance, look_variant):
        """Imprints Ayon metadata onto materials."""
        ayon_materials = self.get_ayon_materials(look_variant)
        if not ayon_materials:
            print(f"[DEBUG v{DEBUG_VERSION}] No materials found to imprint metadata.")
            return
        
        folder_path = instance.get("folderPath", "")
        representation = instance.get("representation", "")
        libpath = instance.get("path", "")
        
        for _, material in ayon_materials:
            if not material:
                continue
            
            material["avalon"] = {
                "schema": "openpype:container-2.0",
                "id": "pyblish.avalon.instance",
                "folderPath": folder_path,
                "task": "shading",
                "variant": look_variant,
                "productType": "look",
                "creator_identifier": "io.openpype.creators.blender.look",
                "productName": instance.get("productName", material.name),
                "active": True,
                "instance_id": str(uuid.uuid4()),
                "creator_attributes": {},
                "publish_attributes": {
                    "ExtractBlend": {"active": True},
                    "IntegrateHeroVersion": {"active": True}
                },
                "loader": "BlendLookLoader",
                "representation": representation,
                "libpath": libpath,
                "objectName": material.name,
            }
        print(f"[DEBUG v{DEBUG_VERSION}] Imprinted Ayon metadata on {len(ayon_materials)} materials.")


    def remove(self, instance):
        """Remove shaderballs, clean up imported meshes, and purge non-existing data."""
        print(f"[DEBUG v{DEBUG_VERSION}] Removing LookFile-related shaderballs and purging data...")
        for collection in bpy.data.collections:
            if collection.name.startswith(AVALON_INSTANCES):
                bpy.data.collections.remove(collection)
        
        shaderballs = [obj for obj in bpy.data.objects if obj.name.startswith(SHADERBALL_OBJECT_NAME)]
        for shaderball in shaderballs:
            bpy.data.objects.remove(shaderball, do_unlink=True)
        
        # Purge unused data including unused materials
        bpy.ops.outliner.orphans_purge(do_local_ids=True, do_linked_ids=True, do_recursive=True)
        print(f"[DEBUG v{DEBUG_VERSION}] Cleanup complete.")

    def get_ayon_materials(self, look_variant):
        """Find all materials linked to meshes under Ayon metadata parent EMPTY objects."""
        collection = bpy.data.collections.get(AVALON_CONTAINERS)
        if not collection:
            print(f"[DEBUG v{DEBUG_VERSION}] AVALON_CONTAINERS collection not found.")
            return []

        ayon_materials = set()
        for parent_obj in collection.objects:
            if parent_obj.type == "EMPTY" and "avalon" in parent_obj:
                asset_name_full = parent_obj["avalon"].get("asset_name", "unknown")
                asset_name = f"{asset_name_full.split('_')[0]}_look{look_variant}"
                
                for child_obj in parent_obj.children:
                    if child_obj.type == "MESH":
                        for mat_slot in child_obj.material_slots:
                            if mat_slot.material:
                                ayon_materials.add((asset_name, mat_slot.material))
        
        print(f"[DEBUG v{DEBUG_VERSION}] Found {len(ayon_materials)} unique materials for variant {look_variant}.")
        return list(ayon_materials)

    def append_shaderball(self):
        """Append the shaderball object from the external .blend file."""
        with bpy.data.libraries.load(str(SHADERBALL_PATH), link=False) as (data_from, data_to):
            if SHADERBALL_OBJECT_NAME in data_from.objects:
                data_to.objects.append(SHADERBALL_OBJECT_NAME)
        return bpy.data.objects.get(SHADERBALL_OBJECT_NAME)

    def ensure_instance_collection(self, asset_name):
        """Ensure a collection exists for shaderballs under AVALON_INSTANCES."""
        if AVALON_INSTANCES not in bpy.data.collections:
            inst_collection = bpy.data.collections.new(AVALON_INSTANCES)
            bpy.context.scene.collection.children.link(inst_collection)
        else:
            inst_collection = bpy.data.collections[AVALON_INSTANCES]
        
        if asset_name not in bpy.data.collections:
            look_collection = bpy.data.collections.new(asset_name)
            inst_collection.children.link(look_collection)
        else:
            look_collection = bpy.data.collections[asset_name]
        
        return look_collection

    def duplicate_shaderballs(self, materials):
        """Duplicate shaderballs while ensuring shared base material data."""
        shaderballs = []
        base_shaderball = self.append_shaderball()
        if not base_shaderball:
            print(f"[DEBUG v{DEBUG_VERSION}] Failed to append shaderball object.")
            return []
        
        for i, (asset_name, material) in enumerate(materials):
            new_shaderball = base_shaderball.copy()
            new_shaderball.data = base_shaderball.data.copy()
            new_shaderball.name = f"look_shaderBall_{asset_name}-{material.name}"
            
            look_collection = self.ensure_instance_collection(asset_name)
            look_collection.objects.link(new_shaderball)
            
            new_shaderball.location.x += i * 2  # Space them out horizontally
            shaderballs.append(new_shaderball)
        
        # Remove duplicate base materials
        for mat in bpy.data.materials:
            if mat.name.startswith("mat_shaderBall_Base") and mat.name.endswith(".001"):
                bpy.data.materials.remove(mat)
        
        return shaderballs

    def replace_placeholder_material(self, shaderballs, materials):
        """Replace the placeholder material on each shaderball."""
        for shaderball, (_, material) in zip(shaderballs, materials):
            for mat_slot in shaderball.material_slots:
                if mat_slot.material and mat_slot.material.name == PLACEHOLDER_MATERIAL_NAME:
                    mat_slot.material = material

    def create_shaderballs_for_look(self, look_variant):
        """Main function to generate shaderballs based on Ayon metadata and selected variant."""
        print(f"[DEBUG v{DEBUG_VERSION}] Generating shaderballs for variant: {look_variant}...")
        ayon_materials = self.get_ayon_materials(look_variant)
        if not ayon_materials:
            print(f"[DEBUG v{DEBUG_VERSION}] No Ayon materials found for variant {look_variant}.")
            return

        shaderballs = self.duplicate_shaderballs(ayon_materials)
        self.replace_placeholder_material(shaderballs, ayon_materials)
        print(f"[DEBUG v{DEBUG_VERSION}] Created {len(shaderballs)} shaderballs for variant {look_variant}.")