import bpy

def get_all_children_recursive(obj, mesh_list):
    """Recursively finds all mesh objects in the hierarchy."""
    for child in obj.children:
        if child.type == 'MESH' and child.data.uv_layers:
            mesh_list.append(child)
        get_all_children_recursive(child, mesh_list)

def setup_furniture_atlas():
    # 1. Gather valid objects: Children of empties starting with "furniture-"

    furniture_meshes = []
    roots = [o for o in bpy.data.objects if o.name.startswith("furniture-")]

    for root in roots:
        # If the root itself is a mesh with UVs, add it
        if root.type == 'MESH' and root.data.uv_layers:
            furniture_meshes.append(root)
        # Find all nested children
        get_all_children_recursive(root, furniture_meshes)

    # Remove duplicates if any (in case of complex parenting)
    furniture_meshes = list(set(furniture_meshes))

    if not furniture_meshes:
        print("No recursive mesh objects found with UVs.")
        return

    furniture_objs = furniture_meshes
    
    if not furniture_objs:
        print("No valid furniture mesh objects found with UVs.")
        return

    # 2. Manage the Atlas Image (FurnitureColor)
    image_name = "FurnitureColor"
    if image_name in bpy.data.images:
        atlas_img = bpy.data.images[image_name]
    else:
        atlas_img = bpy.data.images.new(image_name, width=2048, height=2048)

    # 3. Setup Materials & Nodes
    processed_mats = set()
    for obj in furniture_objs:
        for slot in obj.material_slots:
            mat = slot.material
            if mat and mat not in processed_mats:
                mat.use_nodes = True
                nodes = mat.node_tree.nodes
                
                # Find or create the target node
                bake_node = nodes.get("ATLAS_TARGET")
                if not bake_node:
                    bake_node = nodes.new('ShaderNodeTexImage')
                    bake_node.name = "ATLAS_TARGET"
                
                bake_node.image = atlas_img
                nodes.active = bake_node # Essential for the Bake operator
                processed_mats.add(mat)

    # 4. Selection & Context Setup
    bpy.ops.object.select_all(action='DESELECT')
    for obj in furniture_objs:
        obj.select_set(True)
    bpy.context.view_layer.objects.active = furniture_objs[0]

    # 5. Bake Configuration (Cycles, Diffuse, Color only)
    scene = bpy.context.scene
    scene.render.engine = 'CYCLES'
    scene.cycles.device = 'CPU'
    scene.cycles.samples = 10
    scene.render.bake.use_pass_direct = False
    scene.render.bake.use_pass_indirect = False
    scene.render.bake.use_pass_color = True
    scene.render.bake.target = 'IMAGE_TEXTURES'
    
    print("Baking Furniture Atlas...")
    override = {
        'active_object': bpy.context.view_layer.objects.active,
        'selected_objects': bpy.context.selected_objects,
    }

    with bpy.context.temp_override(**override):
        bpy.ops.object.bake(type='DIFFUSE', margin=2)
    atlas_img.pack()

    # 6. Create and Assign Final Atlas Material
    atlas_mat_name = "M_Furniture_Atlas"
    if atlas_mat_name in bpy.data.materials:
        atlas_mat = bpy.data.materials[atlas_mat_name]
    else:
        atlas_mat = bpy.data.materials.new(name=atlas_mat_name)
        atlas_mat.use_nodes = True
        bsdf = atlas_mat.node_tree.nodes.get("Principled BSDF")
        tex_node = atlas_mat.node_tree.nodes.new('ShaderNodeTexImage')
        tex_node.image = atlas_img
        atlas_mat.node_tree.links.new(tex_node.outputs['Color'], bsdf.inputs['Base Color'])
    nodes = atlas_mat.node_tree.nodes
    bsdf = nodes.get("Principled BSDF")
    if bsdf:
        bsdf.inputs['Roughness'].default_value = 1.0

    # Reassign all valid objects to the new material
    for obj in furniture_objs:
        obj.data.materials.clear()
        obj.data.materials.append(atlas_mat)

    print("Process Complete. Atlas material applied.")
    for mat in bpy.data.materials:
        if mat.users == 0:
            bpy.data.materials.remove(mat)
    bpy.ops.outliner.orphans_purge(do_local_ids=True, do_linked_ids=True, do_recursive=True)

    scene.render.engine = 'BLENDER_EEVEE'
    current_path = bpy.data.filepath
    if current_path:
        new_path = current_path.replace(".blend", "_baked.blend")
        bpy.ops.wm.save_as_mainfile(filepath=new_path)
        print(f"File saved to: {new_path}")
    else:
        print("Warning: File not saved (unsaved blend file).")

setup_furniture_atlas()

