From 4fcf23319f30aa771a32598a1dd9ecbd96fe4674 Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Wed, 21 Aug 2024 11:08:34 +0500
Subject: [PATCH] Show in UI an operator for saving UV

It was lost when we deprecated style UI in materials tab.
Button location - https://i.imgur.com/V57MKU2.png

Also added some minor tweaks.
---
 .../bonsai/bim/module/style/operator.py       | 37 +++++++++++++++----
 src/bonsai/bonsai/bim/module/style/ui.py      |  1 +
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/src/bonsai/bonsai/bim/module/style/operator.py b/src/bonsai/bonsai/bim/module/style/operator.py
index cef47fa6d0..efa0040dde 100644
--- a/src/bonsai/bonsai/bim/module/style/operator.py
+++ b/src/bonsai/bonsai/bim/module/style/operator.py
@@ -882,19 +882,35 @@ class AddSurfaceTexture(bpy.types.Operator):
 class SaveUVToStyle(bpy.types.Operator, tool.Ifc.Operator):
     bl_idname = "bim.save_uv_to_style"
     bl_label = "Save UV To Style"
+    bl_description = "Save active object UV to it's IfcSurfaceStyle's IfcSurfaceTexture"
     bl_options = {"REGISTER", "UNDO"}
 
+    @classmethod
+    def poll(cls, context):
+        if not (obj := context.active_object) or not tool.Ifc.get_entity(obj):
+            cls.poll_message_set("No IFC object selected")
+            return False
+        if not obj.active_material or not tool.Ifc.get_entity(obj.active_material):
+            cls.poll_message_set("Object has no IFC style.")
+            return False
+        return True
+
     def _execute(self, context):
         context = bpy.context
         ifc_file = tool.Ifc.get()
-        material = context.active_object.active_material
         obj = context.active_object
+        assert obj
+
+        # We need IFC material with a texture style since in IFC
+        # UV indexing (IfcTextureCoordinate) requires some IfcSurfaceTexture.
+        material = obj.active_material
+        assert material
 
         # find active style item
-        style_id = material.BIMStyleProperties.ifc_definition_id
-        style = ifc_file.by_id(style_id)
+        style = tool.Ifc.get_entity(material)
+        assert style
 
-        def get_active_representation_items():
+        def get_active_representation_items() -> list[ifcopenshell.entity_instance]:
             representation = tool.Geometry.get_active_representation(obj)
             representation = ifcopenshell.util.representation.resolve_representation(representation)
             items = []
@@ -932,6 +948,14 @@ class SaveUVToStyle(bpy.types.Operator, tool.Ifc.Operator):
             self.report({"ERROR"}, "No texture style found, not yet supported.")
             return {"CANCELLED"}
 
+        mesh = obj.data
+        assert isinstance(mesh, bpy.types.Mesh)
+        uv_layer = mesh.uv_layers.active
+        if not uv_layer:
+            self.report({"ERROR"}, "Object has no UV.")
+            return {"CANCELLED"}
+        uv_layer = uv_layer.uv
+
         # unmap preivously used IfcIndexedTextureMap
         # TODO: remove orphaned data?
         textures = set(texture_style.Textures)
@@ -942,9 +966,8 @@ class SaveUVToStyle(bpy.types.Operator, tool.Ifc.Operator):
             if coord.is_a("IfcIndexedTextureMap") and coord.MappedTo in active_representation_items:
                 coord.Maps = list(set(coord.Maps) - textures)
 
-        mesh = obj.data
         uv_indices = []
-        uv_verts = [uv.vector for uv in mesh.uv_layers.active.uv]
+        uv_verts = [uv.vector for uv in uv_layer]
 
         si_conversion = ifcopenshell.util.unit.calculate_unit_scale(ifc_file)
 
@@ -993,7 +1016,7 @@ class SaveUVToStyle(bpy.types.Operator, tool.Ifc.Operator):
             uv_verts_list.TexCoordsList = uv_verts
             texture_coord.TexCoords = uv_verts_list
 
-        self.report({"INFO"}, f"UV saved to the style {style.Name}")
+        self.report({"INFO"}, f"UV saved to the style '{style.Name}'.")
 
         return {"FINISHED"}
 
diff --git a/src/bonsai/bonsai/bim/module/style/ui.py b/src/bonsai/bonsai/bim/module/style/ui.py
index 6bebbc07cc..8e148476fb 100644
--- a/src/bonsai/bonsai/bim/module/style/ui.py
+++ b/src/bonsai/bonsai/bim/module/style/ui.py
@@ -66,6 +66,7 @@ class BIM_PT_styles(Panel):
             row.operator("bim.select_by_style", text="", icon="RESTRICT_SELECT_OFF").style = style.ifc_definition_id
             op = row.operator("bim.assign_style_to_selected", text="", icon="BRUSH_DATA")
             op.style_id = style.ifc_definition_id
+            op = row.operator("bim.save_uv_to_style", text="", icon="UV")
             op = row.operator("bim.enable_editing_style", text="", icon="GREASEPENCIL")
             op.style = style.ifc_definition_id
             row.operator("bim.remove_style", text="", icon="X").style = style.ifc_definition_id